import time
from machine import Pin, SoftI2C
import json
import network
from simple_mqtt import SimpleMQTT

# -------------------------------------------------------------
# Configuration Loader
# -------------------------------------------------------------
try:
    import config as Config
except ImportError:
    print("ERROR: config.py not found.")
    print("Create config.py on the device with your Wi-Fi and MQTT credentials.")
    while True:
        time.sleep(2)

NODE_NAME = getattr(Config, "NODE_NAME", Config.NODE_ID)
SENSOR_TYPE = getattr(Config, "SENSOR_TYPE", "bme280").lower()
STATE_TOPIC = f"homeassistant/sensor/{Config.NODE_ID}/state"

# Dynamic sensor driver import
if SENSOR_TYPE == "sht41":
    from sht41 import SHT41
elif SENSOR_TYPE == "bme280":
    from bme280 import BME280
else:
    raise ValueError(f"Unsupported SENSOR_TYPE '{SENSOR_TYPE}'. Use 'bme280' or 'sht41'.")


# -------------------------------------------------------------
# Home Assistant MQTT Discovery
# -------------------------------------------------------------
def register_ha_discovery(mqtt, sensor_type):
    device_info = {
        "identifiers": [Config.NODE_ID],
        "name": NODE_NAME,
        "model": f"Pico 2 W ({sensor_type.upper()})",
        "manufacturer": "Raspberry Pi",
    }

    configs = [
        {
            "id": "temperature",
            "name": "Temperature",
            "unit": "°C",
            "cls": "temperature",
            "val_tpl": "{{ value_json.temperature }}",
        },
        {
            "id": "humidity",
            "name": "Humidity",
            "unit": "%",
            "cls": "humidity",
            "val_tpl": "{{ value_json.humidity }}",
        },
    ]

    if sensor_type == "bme280":
        configs.append(
            {
                "id": "pressure",
                "name": "Pressure",
                "unit": "hPa",
                "cls": "atmospheric_pressure",
                "val_tpl": "{{ value_json.pressure }}",
            }
        )

    for c in configs:
        topic = f"homeassistant/sensor/{Config.NODE_ID}/{c['id']}/config"
        payload = {
            "name": c["name"],
            "has_entity_name": True,
            "unique_id": f"{Config.NODE_ID}_{c['id']}",
            "device_class": c["cls"],
            "state_class": "measurement",
            "unit_of_measurement": c["unit"],
            "state_topic": STATE_TOPIC,
            "value_template": c["val_tpl"],
            "device": device_info,
        }

        mqtt.publish(topic, json.dumps(payload), retain=True)
        print(f"Registered HA discovery: {c['id']}")


# -------------------------------------------------------------
# Wi-Fi Connection
# -------------------------------------------------------------
def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print(f"Connecting to Wi-Fi '{Config.WIFI_SSID}'...")
        wlan.connect(Config.WIFI_SSID, Config.WIFI_PASS)
        timeout = 15
        while not wlan.isconnected() and timeout > 0:
            time.sleep(1)
            timeout -= 1

    if not wlan.isconnected():
        raise RuntimeError("Failed to connect to Wi-Fi")
    print(f"Wi-Fi connected. IP: {wlan.ifconfig()[0]}")


# -------------------------------------------------------------
# Panic Handler
# -------------------------------------------------------------
def panic(err=None):
    if err:
        print(f"Fatal error: {err}")
    led = Pin("LED", Pin.OUT)
    while True:
        led.toggle()
        time.sleep(0.1)


# -------------------------------------------------------------
# Main Execution Loop
# -------------------------------------------------------------
def main():
    i2c = SoftI2C(sda=Pin(4), scl=Pin(5), freq=100000)

    if SENSOR_TYPE == "sht41":
        print("Using SHT41 sensor (0x44)")
        sensor = SHT41(i2c)
    else:
        print("Using BME280 sensor (0x77)")
        sensor = BME280(i2c)

    connect_wifi()

    mqtt_user = getattr(Config, "MQTT_USER", None)
    mqtt_pass = getattr(Config, "MQTT_PASS", None)
    mqtt_port = getattr(Config, "MQTT_PORT", 1883)

    mqtt = SimpleMQTT(Config.MQTT_HOST, mqtt_port)
    print(f"Connecting to MQTT broker at {Config.MQTT_HOST}...")
    mqtt.connect(Config.CLIENT_ID, mqtt_user, mqtt_pass)
    print("MQTT connected.")

    register_ha_discovery(mqtt, SENSOR_TYPE)

    interval = getattr(Config, "UPDATE_INTERVAL", 30)
    while True:
        try:
            data = sensor.read_values()

            payload = {k: v for k, v in data.items() if v is not None}
            state_json = json.dumps(payload)

            mqtt.publish(STATE_TOPIC, state_json)
            print(f"[{NODE_NAME}] Published: {state_json}")
        except Exception as e:
            print(f"Error publishing state: {e}")
            try:
                mqtt.disconnect()
                time.sleep(2)
                connect_wifi()
                mqtt.connect(Config.CLIENT_ID, mqtt_user, mqtt_pass)
            except Exception as reconnect_err:
                print(f"Reconnect failed: {reconnect_err}")

        time.sleep(interval)


if __name__ == "__main__":
    time.sleep(2)
    Pin("LED", Pin.OUT).on()

    try:
        main()
    except Exception as e:
        panic(e)