import network
import time
import ujson
from machine import Pin, I2C
from umqtt.simple import MQTTClient
import ssd1306

# =====================================================
# RELAY CONFIG
# =====================================================

RELAY_PINS = [
    23, 19, 18, 5,
    4, 32, 33, 25,
    26, 27, 14, 12,
    13, 15, 2, 0,
    17, 16
]

relays = {
    pin: Pin(pin, Pin.OUT, value=0)
    for pin in RELAY_PINS
}

# =====================================================
# OLED CONFIG
# =====================================================

i2c = I2C(
    0,
    scl=Pin(22),
    sda=Pin(21)
)

oled = ssd1306.SSD1306_I2C(
    128,
    64,
    i2c
)

# =====================================================
# WIFI
# =====================================================

SSID = "Wokwi-GUEST"
PASSWORD = ""

# =====================================================
# HOME ASSISTANT MQTT
# =====================================================

HA_BROKER = "broker.mqttdashboard.com"

HA_TOPIC_SUB = b"home/relay/+/set"

ha_client = MQTTClient(
    client_id="ESP32_HA",
    server=HA_BROKER
)

# =====================================================
# THINGSBOARD MQTT
# =====================================================

TB_BROKER = "thingsboard.cloud"
TB_PORT = 1883

ACCESS_TOKEN = "yBa31inQF6yYOWLf0603"

TB_RPC_TOPIC = b"v1/devices/me/rpc/request/+"
TB_TELEMETRY_TOPIC = "v1/devices/me/telemetry"

tb_client = MQTTClient(
    client_id="ESP32_TB",
    server=TB_BROKER,
    port=TB_PORT,
    user=ACCESS_TOKEN,
    password=""
)

# =====================================================
# OLED DISPLAY
# =====================================================

def update_display(line1="", line2=""):

    oled.fill(0)

    oled.text("HA + TB SYSTEM", 0, 0)

    oled.text(line1, 0, 20)

    oled.text(line2, 0, 35)

    oled.show()

# =====================================================
# WIFI CONNECT
# =====================================================

def connect_wifi():

    update_display("Connecting", "WiFi...")

    wlan = network.WLAN(network.STA_IF)

    wlan.active(True)

    wlan.connect(SSID, PASSWORD)

    while not wlan.isconnected():
        time.sleep(0.2)

    print("WiFi Connected")
    print(wlan.ifconfig())

    update_display("WiFi", "Connected")

# =====================================================
# SEND TELEMETRY TO THINGSBOARD
# =====================================================

def send_tb_telemetry(pin, state):

    try:

        # Dynamic telemetry key
        key = "relay_{}".format(pin)

        data = {
            key: state
        }

        tb_client.publish(
            TB_TELEMETRY_TOPIC,
            ujson.dumps(data)
        )

        print("TB Telemetry:", data)

    except Exception as e:

        print("Telemetry Error:", e)

# =====================================================
# SET RELAY STATE
# =====================================================

def set_relay(pin, state):

    try:

        if pin not in relays:
            print("Invalid GPIO:", pin)
            return

        relays[pin].value(1 if state else 0)

        status_text = "ON" if state else "OFF"

        print("Relay", pin, status_text)

        # OLED
        update_display(
            "GPIO {}".format(pin),
            status_text
        )

        # Publish state back to Home Assistant
        ha_client.publish(
            "home/relay/{}/state".format(pin),
            status_text,
            retain=True
        )

        # Send telemetry to ThingsBoard
        send_tb_telemetry(pin, state)

    except Exception as e:

        print("Relay Error:", e)

# =====================================================
# HOME ASSISTANT CALLBACK
# =====================================================

def ha_callback(topic, msg):

    try:

        print("HA Topic:", topic)
        print("HA Msg:", msg)

        # Topic:
        # home/relay/<pin>/set

        parts = topic.decode().split('/')

        pin = int(parts[2])

        state = True if msg == b'ON' else False

        set_relay(pin, state)

    except Exception as e:

        print("HA Callback Error:", e)

# =====================================================
# THINGSBOARD CALLBACK
# =====================================================

def tb_callback(topic, msg):

    try:

        print("TB Topic:", topic)
        print("TB Msg:", msg)

        data = ujson.loads(msg)

        method = data.get("method")

        params = data.get("params")

        # =========================================
        # GET RELAY STATE
        # =========================================

        if method.startswith("getRelay"):

            pin = int(method.replace("getRelay", ""))

            state = bool(relays[pin].value())

            response_topic = topic.decode().replace(
                "request",
                "response"
            )

            tb_client.publish(
                response_topic,
                ujson.dumps(state)
            )

            print("Returned state:", pin, state)

            return

        # =========================================
        # SINGLE RELAY
        # =========================================

        if method == "setRelay":

            if not isinstance(params, dict):

                print("Invalid params:", params)

                return

            pin = params.get("pin")

            state = params.get("state")

            set_relay(pin, state)

        # =========================================
        # ALL RELAYS
        # =========================================

        elif method == "setAll":

            if not isinstance(params, dict):

                print("Invalid params:", params)

                return

            state = params.get("state")

            for pin in RELAY_PINS:

                set_relay(pin, state)

    except Exception as e:

        print("TB Callback Error:", e)

# =====================================================
# CONNECT WIFI
# =====================================================

connect_wifi()

# =====================================================
# CONNECT HOME ASSISTANT MQTT
# =====================================================

try:

    ha_client.set_callback(ha_callback)

    ha_client.connect()

    ha_client.subscribe(HA_TOPIC_SUB)

    print("Home Assistant MQTT Connected")

except Exception as e:

    print("HA MQTT Error:", e)

# =====================================================
# CONNECT THINGSBOARD MQTT
# =====================================================

try:

    tb_client.set_callback(tb_callback)

    tb_client.connect()

    tb_client.subscribe(TB_RPC_TOPIC)

    print("ThingsBoard Connected")

except Exception as e:

    print("TB MQTT Error:", e)

# =====================================================
# READY
# =====================================================

update_display(
    "System Ready",
    "HA + TB OK"
)

# =====================================================
# MAIN LOOP
# =====================================================

while True:

    try:

        # Listen HA
        ha_client.check_msg()

        # Listen ThingsBoard
        tb_client.check_msg()

        time.sleep(0.1)

    except Exception as e:

        print("Main Loop Error:", e)

        update_display(
            "MQTT Error",
            "Reconnect..."
        )

        time.sleep(5)

        # Reconnect HA
        try:

            ha_client.connect()

            ha_client.subscribe(HA_TOPIC_SUB)

            print("HA Reconnected")

        except Exception as e2:

            print("HA Reconnect Fail:", e2)

        # Reconnect TB
        try:

            tb_client.connect()

            tb_client.subscribe(TB_RPC_TOPIC)

            print("TB Reconnected")

        except Exception as e3:

            print("TB Reconnect Fail:", e3)