import time
import struct
import socket
import network
from machine import Pin, SoftI2C

# -------------------------------------------------------------
# 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)

# Default NODE_NAME to NODE_ID if not explicitly provided in config.py
NODE_NAME = getattr(Config, "NODE_NAME", Config.NODE_ID)
STATE_TOPIC = f"homeassistant/sensor/{Config.NODE_ID}/state"


# -------------------------------------------------------------
# Minimal MQTT 3.1.1 Client
# -------------------------------------------------------------
class SimpleMQTT:
    def __init__(self, host, port=1883):
        self.host = host
        self.port = port
        self.sock = None

    def connect(self, client_id, username=None, password=None):
        addr_info = socket.getaddrinfo(self.host, self.port)
        addr = addr_info[0][-1]
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect(addr)

        # Flags: Clean session (bit 1) + optional credentials
        flags = 0x02
        if username:
            flags |= 0x80
        if password:
            flags |= 0x40

        # Variable header: Protocol Name (MQTT) + Level (4) + Flags + KeepAlive (60s)
        var_header = bytearray([0x00, 0x04, ord('M'), ord('Q'), ord('T'), ord('T'), 0x04, flags, 0x00, 0x3C])

        # Payload
        payload = bytearray()
        payload.extend(self._encode_str(client_id))
        if username:
            payload.extend(self._encode_str(username))
        if password:
            payload.extend(self._encode_str(password))

        body = var_header + payload
        packet = bytearray([0x10]) + self._encode_len(len(body)) + body
        self.sock.write(packet)

        # Read CONNACK (4 bytes: 0x20, 0x02, ack_flags, return_code)
        resp = self.sock.read(4)
        if not resp or resp[0] != 0x20 or resp[3] != 0x00:
            code = resp[3] if resp and len(resp) >= 4 else "nil"
            raise RuntimeError(f"MQTT connection failed: code {code}")

    def publish(self, topic, payload, retain=False):
        cmd = 0x30 | (0x01 if retain else 0x00)
        if isinstance(payload, str):
            payload = payload.encode("utf-8")
        body = self._encode_str(topic) + payload
        packet = bytearray([cmd]) + self._encode_len(len(body)) + body
        self.sock.write(packet)

    def disconnect(self):
        try:
            if self.sock:
                self.sock.write(bytearray([0xE0, 0x00]))
                self.sock.close()
        except Exception:
            pass
        finally:
            self.sock = None

    def _encode_str(self, s):
        raw = s.encode("utf-8")
        return struct.pack("!H", len(raw)) + raw

    def _encode_len(self, length):
        encoded = bytearray()
        while True:
            digit = length % 128
            length //= 128
            if length > 0:
                digit |= 0x80
            encoded.append(digit)
            if length == 0:
                break
        return encoded


# -------------------------------------------------------------
# BME280 Driver (I2C)
# -------------------------------------------------------------
class BME280:
    DEFAULT_ADDR = 0x77

    def __init__(self, i2c, addr=DEFAULT_ADDR):
        self.i2c = i2c
        self.addr = addr

        # Reset sensor
        self.i2c.writeto_mem(self.addr, 0xE0, b"\xB6")
        time.sleep(0.1)

        self._read_calibration_data()

        # Humidity oversampling x1
        self.i2c.writeto_mem(self.addr, 0xF2, b"\x01")
        # Pressure oversampling x1, Temp oversampling x1, Normal mode
        self.i2c.writeto_mem(self.addr, 0xF4, b"\x27")
        # Standby 1000ms, filter off
        self.i2c.writeto_mem(self.addr, 0xF5, b"\xA0")

    def _read_calibration_data(self):
        calib = self.i2c.readfrom_mem(self.addr, 0x88, 24)
        c = struct.unpack("<HhhHhhhhhhhh", calib)
        self.dig_t1, self.dig_t2, self.dig_t3 = c[0], c[1], c[2]
        self.dig_p1, self.dig_p2, self.dig_p3 = c[3], c[4], c[5]
        self.dig_p4, self.dig_p5, self.dig_p6 = c[6], c[7], c[8]
        self.dig_p7, self.dig_p8, self.dig_p9 = c[9], c[10], c[11]

        self.dig_h1 = self.i2c.readfrom_mem(self.addr, 0xA1, 1)[0]
        h = struct.unpack("<hB", self.i2c.readfrom_mem(self.addr, 0xE1, 3)[:3])
        self.dig_h2, self.dig_h3 = h[0], h[1]

        e4 = self.i2c.readfrom_mem(self.addr, 0xE4, 2)
        self.dig_h4 = (e4[0] << 4) | (e4[1] & 0x0F)

        e5_e6 = self.i2c.readfrom_mem(self.addr, 0xE5, 2)
        self.dig_h5 = (e5_e6[0] >> 4) | (e5_e6[1] << 4)

        self.dig_h6 = struct.unpack("<b", self.i2c.readfrom_mem(self.addr, 0xE7, 1))[0]

    def read_values(self):
        raw = self.i2c.readfrom_mem(self.addr, 0xF7, 8)
        raw_p = ((raw[0] << 16) | (raw[1] << 8) | raw[2]) >> 4
        raw_t = ((raw[3] << 16) | (raw[4] << 8) | raw[5]) >> 4
        raw_h = (raw[6] << 8) | raw[7]

        # Temperature calculation (Celsius)
        var1 = (((raw_t >> 3) - (self.dig_t1 << 1)) * self.dig_t2) >> 11
        var2 = (((((raw_t >> 4) - self.dig_t1) * ((raw_t >> 4) - self.dig_t1)) >> 12) * self.dig_t3) >> 14
        t_fine = var1 + var2
        temp = ((t_fine * 5 + 128) >> 8) / 100.0

        # Pressure calculation (hPa)
        p_var1 = t_fine - 128000
        p_var2 = p_var1 * p_var1 * self.dig_p6 + ((p_var1 * self.dig_p5) << 17) + (self.dig_p4 << 35)
        p_var1 = ((p_var1 * p_var1 * self.dig_p3) >> 8) + ((p_var1 * self.dig_p2) << 12)
        p_var1 = (((1 << 47) + p_var1) * self.dig_p1) >> 33

        if p_var1 == 0:
            pressure = 0.0
        else:
            p = 1048576 - raw_p
            p = (((p << 31) - p_var2) * 3125) // p_var1
            p_var1 = (self.dig_p9 * (p >> 13) * (p >> 13)) >> 25
            p_var2 = (self.dig_p8 * p) >> 19
            p = ((p + p_var1 + p_var2) >> 8) + (self.dig_p7 << 4)
            pressure = (p / 256.0) / 100.0

        # Humidity calculation (%)
        h_var = t_fine - 76800
        h_var = (
            ((((raw_h << 14) - (self.dig_h4 << 20) - (self.dig_h5 * h_var)) + 16384) >> 15)
            * (
                (
                    (
                        ((((h_var * self.dig_h6) >> 10) * (((h_var * self.dig_h3) >> 11) + 32768)) >> 10)
                        + 2097152
                    )
                    * self.dig_h2
                    + 8192
                )
                >> 14
            )
        )
        h_var = h_var - (((((h_var >> 15) * (h_var >> 15)) >> 7) * self.dig_h1) >> 4)
        h_var = max(0, min(h_var, 419430400))
        humidity = (h_var >> 12) / 1024.0

        return {
            "temperature": round(temp, 2),
            "pressure": round(pressure, 2),
            "humidity": round(humidity, 2)
        }


# -------------------------------------------------------------
# Home Assistant MQTT Discovery
# -------------------------------------------------------------
def register_ha_discovery(mqtt):
    # Device naming uses NODE_NAME so each board has its own clean device card in HA
    device_info = (
        f'"device":{{'
        f'"identifiers":["{Config.NODE_ID}"],'
        f'"name":"{NODE_NAME}",'
        f'"model":"Pico 2 W",'
        f'"manufacturer":"Raspberry Pi"'
        f'}}'
    )

    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 }}"
        },
        {
            "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"
        unique_id = f"{Config.NODE_ID}_{c['id']}"
        payload = (
            f'{{"name":"{c["name"]}",'
            f'"has_entity_name":true,'
            f'"unique_id":"{unique_id}",'
            f'"device_class":"{c["cls"]}",'
            f'"state_class":"measurement",'
            f'"unit_of_measurement":"{c["unit"]}",'
            f'"state_topic":"{STATE_TOPIC}",'
            f'"value_template":"{c["val_tpl"]}",'
            f'{device_info}}}'
        )
        mqtt.publish(topic, payload, retain=True)
        print(f"Registered HA discovery: {c['id']}")


# -------------------------------------------------------------
# Main Execution Loop
# -------------------------------------------------------------
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]}")


def main():
    # Adjust SDA/SCL pins according to your wiring (Default GP4/GP5 for SoftI2C)
    i2c = SoftI2C(sda=Pin(16), scl=Pin(17), freq=100000)
    bme = 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.")

    # Publish discovery retained messages
    register_ha_discovery(mqtt)

    # Telemetry loop
    interval = getattr(Config, "UPDATE_INTERVAL", 30)
    while True:
        try:
            data = bme.read_values()
            state_json = (
                f'{{"temperature":{data["temperature"]},'
                f'"humidity":{data["humidity"]},'
                f'"pressure":{data["pressure"]}}}'
            )
            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__":
    main()