import gc
import sys
import time
import ujson as json
import uasyncio as asyncio
from machine import Pin, I2C

# 1. Clean memory immediately
gc.collect()

# --- READ DEVICE ID ---
try:
    with open("device_id.txt", "r") as f:
        DEVICE_ID = f.read().strip()
except Exception:
    DEVICE_ID = "ESP32_Inconnu"

# --- GLOBAL APP STATE ---
orchestrator_id = None
cooking_state = None
mqtt_connected = False

# --- MQTT SETUP (Initialized First!) ---
from shared import get_mqtt_client, config, payloads

MQTT_CA_FILE = "/certs/ca.crt"

mqtt_client = get_mqtt_client(
    host="192.168.50.1",
    client_id="smartwave-esp32-demo",
    use_tls=True,
    cafile=MQTT_CA_FILE,
    keepalive=30,
)

# --- HARDWARE & MODULE DEFERRED IMPORTS ---
# We declare variables here, but initialize them AFTER MQTT connects
status_led = None
uart_device = None
mlx_temperature_sensor = None
cookingState = None
log = None
UARTCommand = None
UARTCommandType = None


def init_hardware():
    """Initializes hardware peripherals AFTER MQTT TLS has reserved its memory."""
    global status_led, uart_device, mlx_temperature_sensor
    global cookingState, log, UARTCommand, UARTCommandType
    
    print("[Main] Initializing hardware peripherals...")
    
    # Deferred module imports
    from shared import get_uart, cookingState as cs, logging
    from shared.uart_comm import UARTCommand as UC, UARTCommandType as UCT
    from shared.sensors import RGBLED
    from sensors import temperature_sensor

    cookingState = cs
    log = logging.log
    UARTCommand = UC
    UARTCommandType = UCT

    # Status LED
    status_led = RGBLED(red_pin=21, green_pin=19, blue_pin=18)

    # Hardware UART 2
    uart_device = get_uart(uart_id=2, tx_pin=17, rx_pin=16)

    # I2C Temperature Sensor
    temperature_sensor_i2c = I2C(
        0,
        scl=Pin(25, Pin.IN, Pin.PULL_UP),
        sda=Pin(26, Pin.IN, Pin.PULL_UP),
        freq=100000,
    )
    devices = temperature_sensor_i2c.scan()
    if 0x5A in devices:
        print("[Main] MLX90614 found at address 0x5A!")
    else:
        print("[Main] MLX90614 not found on I2C bus.")
    mlx_temperature_sensor = temperature_sensor.MLX90614(temperature_sensor_i2c)


# --- CALLBACKS ---
def on_cooking_state_change(state):
    if status_led is None:
        return
    print(f"[Main] Cooking state changed to: {state.state}")
    BLINK_INTERVAL_MS = 500
    
    from shared.sensors import RGBLED
    if state.state == cookingState.CookingStates.IDLE:
        status_led.color = RGBLED.OFF
        status_led.blink_off()
    elif state.state == cookingState.CookingStates.COOKING:
        status_led.color = RGBLED.YELLOW
        status_led.blink_off()
    elif state.state == cookingState.CookingStates.STIRRING_REQUIRED:
        status_led.color = RGBLED.ORANGE
        status_led.blink_on(BLINK_INTERVAL_MS)
    elif state.state == cookingState.CookingStates.DONE:
        status_led.color = RGBLED.GREEN
        status_led.blink_off()
    elif state.state == cookingState.CookingStates.ALERT:
        status_led.color = RGBLED.RED
        status_led.blink_on(BLINK_INTERVAL_MS)


def on_received_cooking_state_update(new_state):
    global cooking_state
    if cooking_state is None:
        print("[Main] No active cooking state to update.")
        return
    if log:
        log(f"[Main] Updating cooking state to: {new_state}")
    cooking_state.set_state(new_state)


def on_mqtt_message(message):
    global orchestrator_id, cooking_state
    print("[MQTT] Received message on topic:", message.get("topic"))
    
    payload_data = None
    try:
        payload_data = json.loads(message["payload"])
    except Exception as e:
        print("[MQTT] Payload parsing warning:", e)
    
    topic = message.get("topic")

    # 1. Orchestrator Hello Response
    if (
        topic == config.MQTT_TOPIC_HELLO
        and payload_data
        and payload_data.get("id_microwave") == DEVICE_ID
    ):
        orchestrator_id = payload_data.get("id_orchestrator")
        print("[MQTT] Hello response received from orchestrator:", orchestrator_id)
        try:
            mqtt_client.unsubscribe(config.MQTT_TOPIC_HELLO)
            print("[MQTT] Unsubscribed from topic:", config.MQTT_TOPIC_HELLO)
        except Exception as e:
            print("[MQTT] Unsubscribe error:", e)
            sys.print_exception(e)

    # 2. Cooking Parameters / Sensor Request
    elif (
        topic == config.MQTT_TOPIC_COOKING
        and payload_data
        and payload_data.get("id_microwave") == DEVICE_ID
    ):
        if "cook_time" not in payload_data:
            print("[MQTT] Sensor data requested by orchestrator.")
            obj_temp = mlx_temperature_sensor.read_object_temp() if mlx_temperature_sensor else 0
            amb_temp = mlx_temperature_sensor.read_ambient_temp() if mlx_temperature_sensor else 0
            
            sensor_payload = payloads.mqtt_sensor_data(DEVICE_ID, obj_temp, amb_temp)
            mqtt_client.publish(
                config.MQTT_TOPIC_SENSOR, sensor_payload, qos=config.MQTT_QOS
            )
        else:
            print("[MQTT] Cooking parameters received:", payload_data)
            if cookingState:
                cooking_state = cookingState.CookingState(
                    cook_time=payload_data["cook_time"],
                    power_level=payload_data["power_level"],
                    target_temp=payload_data["target_temp"],
                )
                cooking_state.set_state_change_callback(on_cooking_state_change)
                cooking_state.set_state(cookingState.CookingStates.IDLE)
            
            if uart_device and UARTCommand:
                uart_device.send_as_command(
                    UARTCommand(UARTCommandType.COOKING_PARAMS, payload_data)
                )
                print("[MQTT] Cooking parameters sent to LoRa board.")


mqtt_client.set_callback(on_mqtt_message)


async def connect_mqtt_async():
    """Connects to MQTT safely while memory is clean."""
    global mqtt_connected
    mqtt_connected = False
    
    while True:
        try:
            print("[MQTT] Connecting to broker with TLS...")
            gc.collect()
            mqtt_client.connect()
            print("[MQTT] Connected! Subscribing to topics...")
            mqtt_client.subscribe(config.MQTT_TOPIC_COOKING, qos=config.MQTT_QOS)
            mqtt_client.subscribe(config.MQTT_TOPIC_HELLO, qos=config.MQTT_QOS)
            print("[MQTT] Subscribed successfully!")
            mqtt_connected = True
            return
        except Exception as e:
            print("[MQTT] Connection failed:", e)
            sys.print_exception(e)
            try:
                mqtt_client.close()
            except Exception:
                pass
            print("[MQTT] Retrying connection in 5 seconds...")
            await asyncio.sleep(5)


# --- CONCURRENT ASYNC TASKS ---

async def mqtt_poll_task():
    global mqtt_connected
    last_ping = time.time()

    while True:
        if mqtt_connected:
            try:
                mqtt_client.poll()
                now = time.time()
                if now - last_ping >= 15:
                    if mqtt_client._client:
                        mqtt_client._client.ping()
                    last_ping = now
            except OSError as e:
                print("[MQTT Task] Socket error encountered during poll/ping:", e)
                mqtt_connected = False
                await connect_mqtt_async()

        await asyncio.sleep_ms(30)


async def orchestrator_hello_task():
    global mqtt_connected
    while True:
        if mqtt_connected and orchestrator_id is None:
            print("[Hello Task] Sending initial hello to orchestrator...")
            try:
                mqtt_client.publish(
                    config.MQTT_TOPIC_HELLO,
                    payloads.mqtt_hello(DEVICE_ID),
                    qos=config.MQTT_QOS,
                )
            except OSError as e:
                print("[Hello Task] Hello publish failed:", e)
                mqtt_connected = False
                await connect_mqtt_async()

        await asyncio.sleep(config.MQTT_HELLO_INTERVAL)


async def uart_task():
    while True:
        if uart_device is not None:
            while uart_device.any():
                command = uart_device.read_as_command()
                if command:
                    print(f"[UART Task] Received command: {command.command_type}")
                    if command.command_type == UARTCommandType.COOKING_STATE_UPDATE:
                        new_state = command.payload.get("state", None)
                        print(f"[UART Task] Cooking state update: {new_state}")
                        on_received_cooking_state_update(new_state)
                    else:
                        print(f"[UART Task] Unknown command type: {command.command_type}")
                else:
                    raw_command = uart_device.read()
                    print(f"[UART Task] Received raw command: {raw_command}")

        await asyncio.sleep_ms(20)


async def memory_cleanup_task():
    while True:
        gc.collect()
        await asyncio.sleep(10)


# --- MAIN ENTRY POINT ---
async def main():
    print("[Main] Starting application...")

    # STEP 1: Connect MQTT FIRST (while RAM is unfragmented)
    await connect_mqtt_async()

    # STEP 2: Initialize Hardware & Secondary Modules AFTER connection
    init_hardware()

    # STEP 3: Launch tasks
    asyncio.create_task(mqtt_poll_task())
    asyncio.create_task(orchestrator_hello_task())
    asyncio.create_task(uart_task())
    asyncio.create_task(memory_cleanup_task())

    print("[Main] All tasks running concurrently!")

    while True:
        await asyncio.sleep(3600)


try:
    asyncio.run(main())
except KeyboardInterrupt:
    print("[Main] Program stopped by user.")