"""MQTT protocol collector."""
import json
from typing import Optional

import aiomqtt

from app.collectors.base import BaseCollector


class MqttCollector(BaseCollector):
    """Collect data from devices via MQTT subscription.

    connection_params example:
    {
        "broker": "localhost",
        "port": 1883,
        "topic": "device/INV-001/data",
        "username": "",
        "password": "",
        "data_mapping": {
            "active_power": {"key": "power", "unit": "kW"},
            "voltage": {"key": "voltage", "unit": "V"}
        }
    }
    """

    def __init__(self, device_id, device_code, connection_params, collect_interval=15):
        super().__init__(device_id, device_code, connection_params, collect_interval)
        self._broker = connection_params.get("broker", "localhost")
        self._port = connection_params.get("port", 1883)
        self._topic = connection_params.get("topic", f"device/{device_code}/data")
        self._username = connection_params.get("username", "") or None
        self._password = connection_params.get("password", "") or None
        self._data_mapping = connection_params.get("data_mapping", {})
        self._client: Optional[aiomqtt.Client] = None
        self._latest_data: dict = {}

    async def connect(self):
        # Connection is established in the run loop via context manager
        pass

    async def disconnect(self):
        self._client = None

    async def collect(self) -> dict:
        # Return latest received data; cleared after read
        data = self._latest_data.copy()
        self._latest_data.clear()
        return data

    async def _run(self):
        """Override run loop to use MQTT's push-based model."""
        while self._running:
            try:
                async with aiomqtt.Client(
                    self._broker,
                    port=self._port,
                    username=self._username,
                    password=self._password,
                ) as client:
                    self._client = client
                    self.status = "connected"
                    self.last_error = None
                    self._backoff = 1
                    self.logger.info("MQTT connected to %s:%d", self._broker, self._port)

                    await client.subscribe(self._topic)
                    self.logger.info("Subscribed to %s", self._topic)

                    async for message in client.messages:
                        if not self._running:
                            break
                        try:
                            payload = json.loads(message.payload.decode())
                            data = self._parse_payload(payload)
                            if data:
                                self._latest_data.update(data)
                                await self._save_data(data)
                                from datetime import datetime, timezone
                                self.last_collect_time = datetime.now(timezone.utc)
                        except (json.JSONDecodeError, ValueError) as e:
                            self.logger.warning("Bad MQTT payload on %s: %s", message.topic, e)

            except aiomqtt.MqttError as e:
                self.status = "error"
                self.last_error = str(e)
                self.logger.error("MQTT error for %s: %s", self.device_code, e)
                await self._wait_backoff()
            except Exception as e:
                self.status = "error"
                self.last_error = str(e)
                self.logger.error("Unexpected MQTT error for %s: %s", self.device_code, e)
                await self._wait_backoff()

        self.status = "disconnected"

    def _parse_payload(self, payload: dict) -> dict:
        """Parse MQTT JSON payload into data points.

        If data_mapping is configured, use it. Otherwise, treat all
        numeric top-level keys as data points with empty units.
        """
        data = {}
        if self._data_mapping:
            for data_type, mapping in self._data_mapping.items():
                key = mapping.get("key", data_type)
                unit = mapping.get("unit", "")
                if key in payload:
                    try:
                        data[data_type] = (float(payload[key]), unit)
                    except (TypeError, ValueError):
                        pass
        else:
            for key, value in payload.items():
                if isinstance(value, (int, float)):
                    data[key] = (float(value), "")
        return data