"""
MQTT Connector for Wire.Band
============================

Bridges an MQTT broker into the Wire.Band semantic compression pipeline.

  1. Subscribes to any number of MQTT topic patterns.
  2. Classifies each incoming message to the most specific theta IoT symbol
     (Hardware 0xFC, Metrics 0xF2) using topic path + payload key heuristics.
  3. Encodes the payload as a theta-prefixed compact JSON frame.
  4. Deduplicates semantically identical readings via the semantic cache —
     sensor readings that haven't meaningfully changed never leave the node.
  5. Exposes an async event stream for downstream consumers (HTTP SSE, WebSocket,
     or direct iteration).

Composition with other symbol families:
    0xFC Sensor/IoTProto/Edge  — device-layer semantics (this file)
    0xF2 Metrics/Gauge         — telemetry framing
    0xFB Temporal              — scheduling / duty-cycle metadata
    0xFD Network               — transport (used by the edge client layer)

Dependencies:
    aiomqtt>=2.0.0   (pip install aiomqtt)

Quick start:
    connector = MQTTConnector("mqtt://localhost:1883")
    await connector.connect(["sensors/#", "devices/+/status"])
    async for event in connector.events():
        print(event.topic, event.symbol_glyph, event.compression_ratio)
"""

import asyncio
import json
import logging
import time
import urllib.parse
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Dict, List, Optional, Set

try:
    import aiomqtt
    _AIOMQTT_AVAILABLE = True
except ImportError:
    aiomqtt = None  # type: ignore[assignment]
    _AIOMQTT_AVAILABLE = False

from .symbols import (
    TOPIC_KEYWORD_MAP, PAYLOAD_KEY_MAP, GENERIC_POLL,
    SENSOR_TEMP, SENSOR_VOLTAGE, SENSOR_CURRENT, SENSOR_HUMIDITY,
    SENSOR_PRESSURE, SENSOR_POWER_DRAW, SENSOR_VIBRATION, SENSOR_NOISE,
    SENSOR_ALTITUDE, SENSOR_POLL, SENSOR_THRESHOLD, SENSOR_ALERT,
    GPIO_READ, GPIO_PWM_SET, GPIO_ANALOG_READ, GPIO_ANALOG_WRITE,
    ACTUATOR_MOTOR_START, ACTUATOR_SERVO_SET, ACTUATOR_RELAY_ON,
    ACTUATOR_SOLENOID_OPEN, ACTUATOR_STEPPER_STEP, ACTUATOR_LED_SET, ACTUATOR_BUZZER_ON,
    PROTO_MQTT_PUB, PROTO_MQTT_SUB,
    EDGE_OTA_START, EDGE_WDT_KICK, EDGE_SLEEP, EDGE_WAKE,
    EDGE_TWIN_SYNC, EDGE_TWIN_DELTA, EDGE_INFER, EDGE_PROVISION, EDGE_RESET,
    METRICS_GAUGE_SET, METRICS_GAUGE_CURRENT, METRICS_COUNTER_INC, METRICS_EVENT_EMIT,
)
# Compatibility shims so classifier dict literals still work
class HardwareSensorSymbol:
    TEMP=SENSOR_TEMP; VOLTAGE=SENSOR_VOLTAGE; CURRENT=SENSOR_CURRENT
    HUMIDITY=SENSOR_HUMIDITY; PRESSURE=SENSOR_PRESSURE; POWER_DRAW=SENSOR_POWER_DRAW
    VIBRATION=SENSOR_VIBRATION; NOISE=SENSOR_NOISE; ALTITUDE=SENSOR_ALTITUDE
    POLL=SENSOR_POLL; THRESHOLD=SENSOR_THRESHOLD; ALERT=SENSOR_ALERT

class HardwareGPIOSymbol:
    READ=GPIO_READ; PWM_SET=GPIO_PWM_SET
    ANALOG_READ=GPIO_ANALOG_READ; ANALOG_WRITE=GPIO_ANALOG_WRITE

class HardwareActuatorSymbol:
    MOTOR_START=ACTUATOR_MOTOR_START; SERVO_SET=ACTUATOR_SERVO_SET
    RELAY_ON=ACTUATOR_RELAY_ON; SOLENOID_OPEN=ACTUATOR_SOLENOID_OPEN
    STEPPER_STEP=ACTUATOR_STEPPER_STEP; LED_SET=ACTUATOR_LED_SET; BUZZER_ON=ACTUATOR_BUZZER_ON

class HardwareIoTProtoSymbol:
    MQTT_PUB=PROTO_MQTT_PUB; MQTT_SUB=PROTO_MQTT_SUB

class HardwareEdgeSymbol:
    OTA_START=EDGE_OTA_START; WDT_KICK=EDGE_WDT_KICK; SLEEP=EDGE_SLEEP; WAKE=EDGE_WAKE
    TWIN_SYNC=EDGE_TWIN_SYNC; TWIN_DELTA=EDGE_TWIN_DELTA; EDGE_INFER=EDGE_INFER
    PROVISION=EDGE_PROVISION; RESET=EDGE_RESET

class MetricsGaugeSymbol:
    SET=METRICS_GAUGE_SET; CURRENT=METRICS_GAUGE_CURRENT

class MetricsCounterSymbol:
    INC=METRICS_COUNTER_INC

class MetricsEventSymbol:
    EMIT=METRICS_EVENT_EMIT

class MetricsAlertSymbol:
    pass

HARDWARE_GLYPHS = {}

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Theta-prefixed frame format
#
#   [symbol_high: 1 byte][symbol_low: 1 byte][compact JSON: N bytes]
#
# The 2-byte symbol prefix lets the decoder immediately know the semantic
# category before parsing the payload — enabling domain-specific fast paths.
# ---------------------------------------------------------------------------

FRAME_PREFIX_LEN = 2


@dataclass
class MQTTEvent:
    """A single compressed + classified MQTT event."""

    topic: str
    payload_raw: bytes
    payload_decoded: Any          # parsed JSON dict, list, or str
    symbol: int                   # theta symbol (e.g. 0xFC60 = TEMP)
    symbol_glyph: str             # human glyph (e.g. "S~t")
    compressed_frame: bytes       # [symbol_hi][symbol_lo][compact_json]
    compression_ratio: float      # len(raw) / len(frame)
    ts: float = field(default_factory=time.time)
    cache_hit: bool = False

    @property
    def raw_size(self) -> int:
        return len(self.payload_raw)

    @property
    def compressed_size(self) -> int:
        return len(self.compressed_frame)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "topic":            self.topic,
            "symbol":           hex(self.symbol),
            "symbol_glyph":     self.symbol_glyph,
            "raw_size":         self.raw_size,
            "compressed_size":  self.compressed_size,
            "compression_ratio": round(self.compression_ratio, 2),
            "ts":               self.ts,
            "cache_hit":        self.cache_hit,
        }


# ---------------------------------------------------------------------------
# Topic classifier
# ---------------------------------------------------------------------------

class MQTTTopicClassifier:
    """
    Maps MQTT topic paths + payload keys to theta IoT symbols.

    Resolution order (first match wins):
      1. Exact topic overrides registered via register_topic()
      2. Topic path segment keywords (rightmost segment checked first)
      3. Payload key heuristics (for generic "data" topics)
      4. Fallback: HardwareSensorSymbol.POLL (generic sensor read)

    All keyword matching is case-insensitive.
    """

    # topic keyword → theta symbol
    _TOPIC_KEYWORDS: Dict[str, int] = {
        # Sensors
        "temp":         HardwareSensorSymbol.TEMP,
        "temperature":  HardwareSensorSymbol.TEMP,
        "humidity":     HardwareSensorSymbol.HUMIDITY,
        "pressure":     HardwareSensorSymbol.PRESSURE,
        "voltage":      HardwareSensorSymbol.VOLTAGE,
        "current":      HardwareSensorSymbol.CURRENT,
        "power":        HardwareSensorSymbol.POWER_DRAW,
        "vibration":    HardwareSensorSymbol.VIBRATION,
        "noise":        HardwareSensorSymbol.NOISE,
        "altitude":     HardwareSensorSymbol.ALTITUDE,
        "sensor":       HardwareSensorSymbol.POLL,
        "threshold":    HardwareSensorSymbol.THRESHOLD,
        "alert":        HardwareSensorSymbol.ALERT,
        # Actuators
        "motor":        HardwareActuatorSymbol.MOTOR_START,
        "servo":        HardwareActuatorSymbol.SERVO_SET,
        "relay":        HardwareActuatorSymbol.RELAY_ON,
        "solenoid":     HardwareActuatorSymbol.SOLENOID_OPEN,
        "stepper":      HardwareActuatorSymbol.STEPPER_STEP,
        "led":          HardwareActuatorSymbol.LED_SET,
        "buzzer":       HardwareActuatorSymbol.BUZZER_ON,
        "actuator":     HardwareActuatorSymbol.MOTOR_START,
        # GPIO
        "gpio":         HardwareGPIOSymbol.READ,
        "pin":          HardwareGPIOSymbol.READ,
        "pwm":          HardwareGPIOSymbol.PWM_SET,
        "analog":       HardwareGPIOSymbol.ANALOG_READ,
        "adc":          HardwareGPIOSymbol.ANALOG_READ,
        "dac":          HardwareGPIOSymbol.ANALOG_WRITE,
        # Edge lifecycle
        "ota":          HardwareEdgeSymbol.OTA_START,
        "update":       HardwareEdgeSymbol.OTA_START,
        "watchdog":     HardwareEdgeSymbol.WDT_KICK,
        "wdt":          HardwareEdgeSymbol.WDT_KICK,
        "sleep":        HardwareEdgeSymbol.SLEEP,
        "wake":         HardwareEdgeSymbol.WAKE,
        "twin":         HardwareEdgeSymbol.TWIN_SYNC,
        "shadow":       HardwareEdgeSymbol.TWIN_SYNC,
        "provision":    HardwareEdgeSymbol.PROVISION,
        "reset":        HardwareEdgeSymbol.RESET,
        "inference":    HardwareEdgeSymbol.EDGE_INFER,
        # Metrics / telemetry
        "metrics":      MetricsGaugeSymbol.SET,
        "telemetry":    MetricsGaugeSymbol.SET,
        "counter":      MetricsCounterSymbol.INC,
        "event":        MetricsEventSymbol.EMIT,
        "log":          MetricsEventSymbol.EMIT,
        "status":       MetricsGaugeSymbol.CURRENT,
        # Protocol commands
        "cmd":          HardwareIoTProtoSymbol.MQTT_PUB,
        "command":      HardwareIoTProtoSymbol.MQTT_PUB,
    }

    # payload key → theta symbol (used when topic is generic, e.g. "data")
    _PAYLOAD_KEYS: Dict[str, int] = {
        "temperature":  HardwareSensorSymbol.TEMP,
        "temp":         HardwareSensorSymbol.TEMP,
        "humidity":     HardwareSensorSymbol.HUMIDITY,
        "pressure":     HardwareSensorSymbol.PRESSURE,
        "voltage":      HardwareSensorSymbol.VOLTAGE,
        "current":      HardwareSensorSymbol.CURRENT,
        "rpm":          HardwareActuatorSymbol.MOTOR_START,
        "speed":        HardwareActuatorSymbol.MOTOR_START,
        "angle":        HardwareActuatorSymbol.SERVO_SET,
        "state":        MetricsGaugeSymbol.CURRENT,
        "value":        MetricsGaugeSymbol.SET,
        "count":        MetricsCounterSymbol.INC,
        "level":        MetricsGaugeSymbol.SET,
    }

    def __init__(self) -> None:
        self._exact: Dict[str, int] = {}

    def classify(self, topic: str, payload: Any) -> int:
        """Return the best theta symbol for this message."""
        # 1. Exact topic override
        if topic in self._exact:
            return self._exact[topic]

        # 2. Topic path segments (rightmost first = most specific).
        # Longest-match wins — deterministic and prevents short keywords
        # ("pin", "log") from incorrectly firing on longer segments
        # ("spindle", "analog").
        parts = topic.lower().replace("-", "_").split("/")
        for part in reversed(parts):
            matches = [(kw, sym) for kw, sym in self._TOPIC_KEYWORDS.items() if kw in part]
            if matches:
                return max(matches, key=lambda x: len(x[0]))[1]

        # 3. Payload key heuristics
        if isinstance(payload, dict):
            for key in payload:
                sym = self._PAYLOAD_KEYS.get(key.lower())
                if sym is not None:
                    return sym

        # 4. Fallback
        return HardwareSensorSymbol.POLL

    def register_topic(self, topic: str, symbol: int) -> None:
        """Pin an exact MQTT topic to a specific theta symbol."""
        self._exact[topic] = symbol

    def register_keyword(self, keyword: str, symbol: int) -> None:
        """Add a topic keyword → symbol mapping."""
        self._TOPIC_KEYWORDS[keyword.lower()] = symbol


# ---------------------------------------------------------------------------
# MQTT Connector
# ---------------------------------------------------------------------------

class MQTTConnector:
    """
    Bridges an MQTT broker to the Wire.Band compression pipeline.

    Args:
        broker_url:     MQTT broker URL, e.g. "mqtt://localhost:1883" or
                        "mqtts://broker.example.com:8883"
        client_id:      MQTT client identifier (default: wire-band-connector)
        username:       Optional broker username
        password:       Optional broker password
        keepalive:      MQTT keepalive interval in seconds (default: 60)
        max_queue_size: Max buffered events before backpressure drop (default: 10,000)
        cache:          Optional SemanticInferenceCache for deduplication.
                        When set, semantically identical readings that arrive
                        within the TTL window are deduplicated at the edge.
        delta_threshold: Minimum fractional change required to pass a sensor
                         reading through (0.0 = pass all, 0.02 = 2% change).
                         Only applies to numeric scalar payloads.

    Usage:
        async with MQTTConnector("mqtt://localhost:1883") as conn:
            await conn.subscribe(["sensors/#", "devices/+/twin"])
            async for event in conn.events():
                print(f"{event.topic} → {event.symbol_glyph} "
                      f"({event.compression_ratio:.1f}x)")
    """

    def __init__(
        self,
        broker_url: str = "mqtt://localhost:1883",
        client_id: str = "wire-band-connector",
        username: Optional[str] = None,
        password: Optional[str] = None,
        keepalive: int = 60,
        max_queue_size: int = 10_000,
        cache: Optional[Any] = None,
        delta_threshold: float = 0.0,
    ) -> None:
        self._broker_url = broker_url
        self._client_id = client_id
        self._username = username
        self._password = password
        self._keepalive = keepalive
        self._cache = cache
        self._delta_threshold = delta_threshold

        self._queue: asyncio.Queue = asyncio.Queue(maxsize=max_queue_size)
        self._classifier = MQTTTopicClassifier()
        self._subscriptions: Set[str] = set()
        self._connected = False
        self._task: Optional[asyncio.Task] = None

        # Last-value store for delta filtering (topic → last numeric value)
        self._last_values: Dict[str, float] = {}

        self._stats: Dict[str, Any] = {
            "messages_received":    0,
            "messages_compressed":  0,
            "messages_deduplicated": 0,
            "messages_delta_filtered": 0,
            "bytes_in":             0,
            "bytes_out":            0,
            "errors":               0,
            "started_at":           time.time(),
        }

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def connect(self, topics: List[str]) -> None:
        """Connect to the broker and subscribe to *topics*."""
        if not _AIOMQTT_AVAILABLE:
            raise ImportError(
                "aiomqtt is required for the MQTT connector.\n"
                "Install with: pip install aiomqtt>=2.0.0"
            )
        self._subscriptions = set(topics)
        self._task = asyncio.create_task(self._run(topics))
        # Give the connection a moment to establish before returning
        for _ in range(20):
            if self._connected:
                break
            await asyncio.sleep(0.05)
        logger.info(
            "MQTTConnector connected to %s, subscribed to %s",
            self._broker_url, topics
        )

    async def subscribe(self, topics: List[str]) -> None:
        """Alias for connect() — subscribe to additional topics."""
        await self.connect(topics)

    async def disconnect(self) -> None:
        """Gracefully disconnect from the broker."""
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        self._connected = False
        logger.info("MQTTConnector disconnected from %s", self._broker_url)

    async def events(self) -> AsyncIterator[MQTTEvent]:
        """Async generator yielding compressed MQTTEvents as they arrive."""
        while True:
            event = await self._queue.get()
            yield event

    def register_topic_symbol(self, topic: str, symbol: int) -> None:
        """Pin an exact topic to a theta symbol."""
        self._classifier.register_topic(topic, symbol)

    def register_keyword_symbol(self, keyword: str, symbol: int) -> None:
        """Add a topic keyword → symbol mapping."""
        self._classifier.register_keyword(keyword, symbol)

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def subscriptions(self) -> Set[str]:
        return set(self._subscriptions)

    def stats(self) -> Dict[str, Any]:
        s = dict(self._stats)
        uptime = time.time() - s.pop("started_at")
        bytes_in = s["bytes_in"] or 1
        bytes_out = s["bytes_out"] or 1
        return {
            **s,
            "uptime_seconds":     round(uptime, 1),
            "compression_ratio":  round(bytes_in / bytes_out, 2),
            "connected":          self._connected,
            "subscriptions":      list(self._subscriptions),
        }

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    async def __aenter__(self) -> "MQTTConnector":
        return self

    async def __aexit__(self, *_: Any) -> None:
        await self.disconnect()

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    async def _run(self, topics: List[str]) -> None:
        parsed = urllib.parse.urlparse(self._broker_url)
        hostname = parsed.hostname or "localhost"
        port = parsed.port or (8883 if parsed.scheme == "mqtts" else 1883)
        tls = parsed.scheme == "mqtts"

        try:
            client_kwargs: Dict[str, Any] = dict(
                hostname=hostname,
                port=port,
                identifier=self._client_id,
                keepalive=self._keepalive,
            )
            if self._username:
                client_kwargs["username"] = self._username
            if self._password:
                client_kwargs["password"] = self._password
            if tls:
                import ssl
                client_kwargs["tls_context"] = ssl.create_default_context()

            async with aiomqtt.Client(**client_kwargs) as client:
                self._connected = True
                for topic in topics:
                    await client.subscribe(topic)

                async for message in client.messages:
                    await self._process(str(message.topic), bytes(message.payload))

        except asyncio.CancelledError:
            raise
        except Exception as exc:
            self._stats["errors"] += 1
            self._connected = False
            logger.error("MQTTConnector error: %s", exc)

    async def _process(self, topic: str, raw: bytes) -> None:
        self._stats["messages_received"] += 1
        self._stats["bytes_in"] += len(raw)

        try:
            # Decode payload
            try:
                payload = json.loads(raw)
            except (json.JSONDecodeError, UnicodeDecodeError):
                payload = raw.decode("utf-8", errors="replace")

            # Delta filter — skip if numeric value hasn't changed enough
            if self._delta_threshold > 0.0 and self._should_filter(topic, payload):
                self._stats["messages_delta_filtered"] += 1
                return

            # Semantic deduplication via cache
            cache_hit = False
            if self._cache is not None:
                cache_key = self._cache_key(topic, payload)
                cached = await self._cache.get(cache_key)
                if cached:
                    self._stats["messages_deduplicated"] += 1
                    cache_hit = True

            # Classify to theta symbol
            symbol = self._classifier.classify(topic, payload)
            glyph = HARDWARE_GLYPHS.get(symbol, hex(symbol))

            # Encode as theta-prefixed compact frame
            frame = self._encode(symbol, topic, payload)
            ratio = len(raw) / max(len(frame), 1)

            self._stats["messages_compressed"] += 1
            self._stats["bytes_out"] += len(frame)

            # Store in cache on miss
            if self._cache is not None and not cache_hit:
                await self._cache.put(self._cache_key(topic, payload), frame)

            event = MQTTEvent(
                topic=topic,
                payload_raw=raw,
                payload_decoded=payload,
                symbol=symbol,
                symbol_glyph=glyph,
                compressed_frame=frame,
                compression_ratio=ratio,
                cache_hit=cache_hit,
            )

            try:
                self._queue.put_nowait(event)
            except asyncio.QueueFull:
                logger.warning("Event queue full — dropping message on %s", topic)

        except Exception as exc:
            self._stats["errors"] += 1
            logger.error("MQTTConnector._process error [%s]: %s", topic, exc)

    def _encode(self, symbol: int, topic: str, payload: Any) -> bytes:
        """
        Theta-prefixed compact JSON frame:
            [symbol_hi: 1B][symbol_lo: 1B][compact_json: NB]

        The 2-byte prefix classifies the semantic category so decoders
        can apply domain-specific fast paths without parsing the JSON.
        The compact JSON eliminates all whitespace for additional savings.
        """
        body = json.dumps(
            {"t": topic, "d": payload},
            separators=(",", ":"),
            ensure_ascii=False,
        ).encode("utf-8")
        prefix = bytes([symbol >> 8, symbol & 0xFF])
        return prefix + body

    def _cache_key(self, topic: str, payload: Any) -> str:
        if isinstance(payload, dict):
            return f"mqtt:{topic}:{json.dumps(payload, sort_keys=True, separators=(',', ':'))}"
        return f"mqtt:{topic}:{payload}"

    def _should_filter(self, topic: str, payload: Any) -> bool:
        """Return True if the value hasn't changed enough to forward."""
        value: Optional[float] = None
        if isinstance(payload, (int, float)):
            value = float(payload)
        elif isinstance(payload, dict) and len(payload) == 1:
            v = next(iter(payload.values()))
            if isinstance(v, (int, float)):
                value = float(v)

        if value is None:
            return False  # non-scalar — always pass through

        last = self._last_values.get(topic)
        if last is None:
            self._last_values[topic] = value
            return False

        if last == 0.0:
            changed = abs(value) > self._delta_threshold
        else:
            changed = abs(value - last) / abs(last) >= self._delta_threshold

        if changed:
            self._last_values[topic] = value
        return not changed
