"""Small MQTT compatibility layer for CPython and MicroPython.

The wrapper keeps the broker host explicit so embedded clients can point to a
real IP address instead of localhost.
"""

try:
    import paho.mqtt.client as _mqtt
    BACKEND_NAME = "paho"
    IS_MICROPYTHON = False
except ImportError:
    try:
        from umqtt.simple import MQTTClient as _MQTTClient
        import _thread
        import gc
        BACKEND_NAME = "umqtt.simple"
        IS_MICROPYTHON = True
    except ImportError as exc:
        raise ImportError("No MQTT client found. Expected paho.mqtt or umqtt.") from exc


DEFAULT_PORT = 8884
DEFAULT_CA_FILE = "orchestrateur/mqtt/certs/ca.crt"


def _resolve_host(host):
    if host:
        return host

    try:
        import os

        getenv = getattr(os, "getenv", None)
        if getenv is not None:
            host = getenv("MQTT_BROKER_HOST")
    except Exception:
        host = None

    if not host:
        raise ValueError("MQTT broker host is required. Pass the broker IP address instead of localhost.")

    return host


def _ensure_bytes(payload):
    if payload is None:
        return b""
    if isinstance(payload, bytes):
        return payload
    if isinstance(payload, bytearray):
        return bytes(payload)
    return str(payload).encode()


def _read_file_bytes(path):
    with open(path, "rb") as handle:
        return handle.read()


class BrokerClient:
    """Small MQTT client with a normalised API across runtimes."""

    def __init__(self, host=None, port=DEFAULT_PORT, client_id=None, use_tls=True, cafile=None, certfile=None, keyfile=None, ssl_params=None, tls_insecure=False, username=None, password=None, keepalive=60):
        self.host = _resolve_host(host)
        self.port = port
        self.client_id = client_id
        self.use_tls = use_tls
        self.cafile = cafile or DEFAULT_CA_FILE
        self.certfile = certfile
        self.keyfile = keyfile
        self.ssl_params = ssl_params
        self.tls_insecure = tls_insecure
        self.username = username
        self.password = password
        self.keepalive = keepalive
        self._client = None
        self._callback = None
        self._messages = []
        self._cadata = None  # Cache cert bytes to prevent heap fragmentation
        
        # Thread safety lock for MicroPython socket reads/writes
        if IS_MICROPYTHON:
            self._lock = _thread.allocate_lock()

    def set_callback(self, callback):
        self._callback = callback
        if self._client is not None and not IS_MICROPYTHON:
            self._client.on_message = self._on_message

    def _store_message(self, topic, payload, qos=None, retain=False):
        message = {
            "topic": topic,
            "payload": payload,
            "qos": qos,
            "retain": retain,
        }
        self._messages.append(message)
        if self._callback is not None:
            self._callback(message)

    def _on_message(self, client, userdata, msg):
        self._store_message(msg.topic, msg.payload, getattr(msg, "qos", None), getattr(msg, "retain", False))

    def open(self):
        if self._client is not None:
            return self._client

        if IS_MICROPYTHON:
            gc.collect()  # Clean Python heap before importing/allocating SSL
            import ssl
            ssl_params = self.ssl_params
            
            if self.use_tls and ssl_params is None:
                # OPTION A: If broker uses 'require_certificate false' and self-signed certs:
                # Do NOT pass cadata when cert_reqs is CERT_NONE to save ~20KB of C-DRAM
                ssl_params = {
                    "cert_reqs": ssl.CERT_NONE,
                    "server_hostname": self.host
                }

                # OPTION B: If strict CA validation IS required, load cadata ONLY with CERT_REQUIRED:
                # ssl_params = {
                #     "cert_reqs": ssl.CERT_REQUIRED,
                #     "cadata": _read_file_bytes(self.cafile),
                #     "server_hostname": self.host
                # }
                
            client = _MQTTClient(
                self.client_id or "smartWave-client",
                self.host,
                port=self.port,
                user=self.username,
                password=self.password,
                keepalive=self.keepalive,
                ssl=self.use_tls,
                ssl_params=ssl_params,
            )
            self._client = client
            return self._client

        client = _mqtt.Client(client_id=self.client_id or "", clean_session=True, protocol=4, transport="tcp")
        if self.username is not None or self.password is not None:
            client.username_pw_set(self.username, self.password)
        if self.use_tls:
            tls_kwargs = {}
            if self.cafile is not None:
                tls_kwargs["ca_certs"] = self.cafile
            if self.certfile is not None:
                tls_kwargs["certfile"] = self.certfile
            if self.keyfile is not None:
                tls_kwargs["keyfile"] = self.keyfile
            if tls_kwargs:
                client.tls_set(**tls_kwargs)
            else:
                client.tls_set()
            if self.tls_insecure:
                client.tls_insecure_set(True)
        client.on_message = self._on_message
        self._client = client
        return self._client

    def connect(self):
        if IS_MICROPYTHON:
            gc.collect()  # Force C & Python memory cleanup right before TLS handshake

        if self._client is not None:
            self.close()
        
        client = self.open()
        
        try:
            if IS_MICROPYTHON:
                gc.collect()  # Sweep memory right before umqtt calls ssl.wrap_socket()
                with self._lock:
                    client.connect()
                return client

            client.connect(self.host, self.port, self.keepalive)
            return client
        except Exception as e:
            print("MQTT connection failed, closing client and releasing memory.")
            print("Exception:", e)
            self.close()
            raise

    def publish(self, topic, payload, qos=2, retain=False):
        client = self.open()
        payload_bytes = _ensure_bytes(payload)
        if IS_MICROPYTHON:
            with self._lock:
                return client.publish(topic, payload_bytes, retain=retain, qos=qos)
        
        if isinstance(topic, bytes):
            topic = topic.decode('utf-8')
            
        return client.publish(topic, payload_bytes, qos=qos, retain=retain)

    def subscribe(self, topic, qos=2):
        client = self.open()
        if IS_MICROPYTHON:
            with self._lock:
                client.set_callback(self._on_micropython_message)
                return client.subscribe(topic, qos=qos)
            
        if isinstance(topic, bytes):
            topic = topic.decode('utf-8')
            
        return client.subscribe(topic, qos=qos)
    
    def unsubscribe(self, topic):
        client = self.open()
        if IS_MICROPYTHON:
            import struct
            import time
            topic_bytes = topic if isinstance(topic, bytes) else topic.encode('utf-8')
            
            # 1. Increment and lock the PID for THIS specific request
            client.pid = (client.pid % 65535) + 1
            sent_pid = client.pid  # <-- Store local copy
            
            # 2. Construct UNSUBSCRIBE packet
            rem_len = 2 + 2 + len(topic_bytes)
            pkt = bytearray(b"\xa2\0\0\0")
            struct.pack_into("!BH", pkt, 1, rem_len, sent_pid)
            
            # 3. Write packet to socket
            client.sock.write(pkt)
            client._send_str(topic_bytes)
            
            # 4. Wait for UNSUBACK (0xB0)
            start = time.time()
            while time.time() - start < 3:
                op = client.wait_msg()
                if op == 0xB0:
                    resp = bytearray(3)
                    read_bytes = 0
                    while read_bytes < 3:
                        chunk = client.sock.read(3 - read_bytes)
                        if chunk:
                            resp[read_bytes:read_bytes + len(chunk)] = chunk
                            read_bytes += len(chunk)
                        else:
                            time.sleep_ms(10)
                    
                    # Compare against sent_pid instead of client.pid
                    resp_pid = (resp[1] << 8) | resp[2]
                    if resp_pid != sent_pid:
                        print(f"[MQTT] UNSUBACK PID mismatch (expected {sent_pid}, got {resp_pid})")
                    return client
            
            return client
            
        if isinstance(topic, bytes):
            topic = topic.decode('utf-8')
            
        return client.unsubscribe(topic)

    def _on_micropython_message(self, topic, payload):
        self._store_message(topic, payload, None, False)

    def poll(self, timeout=0.1):
        if self._client is None:
            return None
        if IS_MICROPYTHON:
            with self._lock:
                return self._client.check_msg()
        return self._client.loop(timeout=timeout)

    def wait(self):
        if self._client is None:
            return None
        if IS_MICROPYTHON:
            with self._lock:
                return self._client.wait_msg()
        return self._client.loop_forever()

    def get_message(self):
        if not self._messages:
            return None
        return self._messages.pop(0)

    def close(self):
        """Safely clean up socket context without causing ESP32 C panics."""
        if self._client is None:
            return
            
        if IS_MICROPYTHON:
            with self._lock:
                try:
                    if hasattr(self._client, "sock") and self._client.sock:
                        self._client.sock.close()
                except Exception:
                    pass
                finally:
                    if hasattr(self._client, "sock"):
                        self._client.sock = None
                    self._client = None
                    gc.collect()  # Immediately reclaim freed socket & mbedTLS RAM
        else:
            try:
                self._client.disconnect()
            except Exception:
                pass
            finally:
                self._client = None
                
    def ping(self):
        """Thread-safe PINGREQ wrapper for MicroPython."""
        if self._client is None:
            return
        if IS_MICROPYTHON:
            with self._lock:
                return self._client.ping()
        else:
            # Paho handles keepalives automatically via loop_start/loop
            pass

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc, traceback):
        self.close()


def connect(host=None, **client_kwargs):
    return BrokerClient(host=host, **client_kwargs)


def publish(host, topic, payload, **client_kwargs):
    qos = client_kwargs.pop("qos", 2)
    retain = client_kwargs.pop("retain", False)
    client = connect(host=host, **client_kwargs)
    client.connect()
    try:
        return client.publish(topic, payload, qos=qos, retain=retain)
    finally:
        client.close()