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

# 1. Clean memory immediately before performing any operations
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
unsubscribed_hello = False

# --- ASYNC SIGNALS & QUEUES ---
# Event to signal when orchestrator requests sensor data (prevents MQTT lock deadlock)
sensor_request_event = None

# --- MQTT SETUP ---
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 ---
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 RAM."""
    global status_led, uart_device, mlx_temperature_sensor
    global cookingState, log, UARTCommand, UARTCommandType

    print("[Main] Initializing hardware peripherals...")

    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 = RGBLED(red_pin=21, green_pin=19, blue_pin=18)
    uart_device = get_uart(uart_id=2, tx_pin=17, rx_pin=16)

    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)


def on_received_cooking_state_update(state, is_error=False, is_terminated=False):
    """Callback executed when state changes are received from the LoRa board over UART."""
    if cooking_state:
        if is_error:
            cooking_state.set_state(cookingState.CookingStates.ERROR)
        elif is_terminated:
            cooking_state.set_state(cookingState.CookingStates.ABORTED)
        else:
            cooking_state.set_state(state)


def on_cooking_state_change(state):
    """Callback executed whenever local cooking state transitions."""
    if status_led and cookingState:
        if state == cookingState.CookingStates.IDLE:
            status_led.set_color(0, 0, 0)  # Off
        elif state == cookingState.CookingStates.PREHEATING:
            status_led.set_color(255, 165, 0)  # Orange
        elif state == cookingState.CookingStates.COOKING:
            status_led.set_color(255, 0, 0)  # Red
        elif state == cookingState.CookingStates.DONE:
            status_led.set_color(0, 255, 0)  # Green
        elif state in (
            cookingState.CookingStates.ERROR,
            cookingState.CookingStates.ABORTED,
        ):
            status_led.set_color(255, 0, 255)  # Magenta/Purple


def on_mqtt_message(message):
    """Sync callback: Lightweight! Only updates variables or triggers async signals."""
    global orchestrator_id, cooking_state, unsubscribed_hello
    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)
        if not unsubscribed_hello:
            unsubscribed_hello = True
            try:
                mqtt_client.unsubscribe(config.MQTT_TOPIC_HELLO)
                print("[MQTT] Successfully unsubscribed from topic:", config.MQTT_TOPIC_HELLO)
            except Exception as e:
                print("[MQTT] Unsubscribe error:", 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! Triggering async publisher...")
            # Trigger async event instead of calling publish() directly inside lock context!
            sensor_request_event.set()
        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 over UART.")


# --- DEDICATED ASYNC TASK FOR SENSOR PUBLISHING ---
async def sensor_publisher_task():
    """Waits for sensor_request_event, reads hardware, and publishes outside the MQTT lock."""
    while True:
        await sensor_request_event.wait()
        sensor_request_event.clear()

        print("[Sensor Task] Reading temperature sensors...")
        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)

        try:
            print("[Sensor Task] Publishing sensor data to MQTT...")
            mqtt_client.publish(
                config.MQTT_TOPIC_SENSOR, sensor_payload, qos=config.MQTT_QOS
            )
            print("[Sensor Task] Sensor data successfully published:", sensor_payload)
        except Exception as e:
            print("[Sensor Task] Failed to publish sensor data:", e)


async def uart_task():
    """Polls incoming UART messages from the LoRa board using dynamic method fallback."""
    while True:
        if uart_device:
            try:
                cmd = uart_device.read_as_command()

                if cmd:
                    print("[UART] Command received from LoRa board:", cmd)
                    if (
                        hasattr(cmd, "command_type")
                        and cmd.command_type == UARTCommandType.STATE_UPDATE
                        and on_received_cooking_state_update
                    ):
                        on_received_cooking_state_update(
                            cmd.payload.get("state"),
                            cmd.payload.get("is_error", False),
                            cmd.payload.get("is_terminated", False),
                        )
            except Exception as e:
                print("[UART Task] Error reading command:", e)

        await asyncio.sleep_ms(50)


async def connect_mqtt_async():
    global mqtt_connected, mqtt_client
    mqtt_connected = False

    while True:
        try:
            print("[MQTT] Connecting to broker with TLS...")
            # Re-instantiate client to clear old socket buffers
            gc.collect()
            mqtt_client = get_mqtt_client(
                host="192.168.50.1", # TODO : Use config.MQTT_BROKER_HOST instead of hardcoding
                port=8884,
                client_id="smartwave-esp32-demo",
                use_tls=True,
                cafile=MQTT_CA_FILE,
                keepalive=30,
            )
            mqtt_client.set_callback(on_mqtt_message)
            
            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
            
            # Force heap cleanup before sleeping
            del mqtt_client
            gc.collect()
            print(f"[MQTT] Free RAM after cleanup: {gc.mem_free()} bytes")
            print("[MQTT] Retrying connection in 5 seconds...")
            await asyncio.sleep(5)

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 orchestrator_id is not None:
            # Hello successfully acknowledged! Stop looping this task.
            print("[Hello Task] Orchestrator acknowledged. Stopping hello task.")
            break
                
        if mqtt_connected:
            print("[Hello Task] Sending initial hello to orchestrator...")
            try:
                if mqtt_client is None:
                    print("[Hello Task] MQTT client is None. Attempting to reconnect...")
                    await connect_mqtt_async()
                    
                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 asyncio.sleep(config.MQTT_HELLO_INTERVAL)


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


# --- MAIN ENTRY POINT ---
async def main():
    global sensor_request_event
    print("[Main] Starting application...")
    
    # Initialize loop-bound events
    sensor_request_event = asyncio.Event()

    await connect_mqtt_async()
    init_hardware()

    # Launch background tasks
    asyncio.create_task(mqtt_poll_task())
    asyncio.create_task(orchestrator_hello_task())
    asyncio.create_task(sensor_publisher_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.")