"""
Demeter Server — CoAP Observer Client

Wraps aiocoap to manage Observe subscriptions to ESP sensor nodes.
Each subscription receives push notifications when sensor values change,
updating the in-memory DeviceStore.

Also provides one-shot CoAP GET/PUT for the REST-to-CoAP bridge.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import TYPE_CHECKING, Any, Optional

import aiocoap
from aiocoap import Code, Message

if TYPE_CHECKING:
    from .device_store import DeviceStore
    from .metrics import MetricsCollector

logger = logging.getLogger(__name__)


class CoapObserverClient:
    """
    Manages CoAP Observe subscriptions and one-shot requests.

    Lifecycle:
        client = CoapObserverClient(store, metrics)
        await client.startup()
        await client.subscribe("esp32-plant-01", "192.168.1.100", 5683, "sensors/temperature")
        ...
        await client.shutdown()
    """

    def __init__(
        self,
        store: DeviceStore,
        metrics: Optional[MetricsCollector] = None,
        request_timeout: float = 10.0,
        reconnect_base: float = 5.0,
        reconnect_max: float = 60.0,
    ):
        self._store = store
        self._metrics = metrics
        self._context: Optional[aiocoap.Context] = None
        self._subscriptions: dict[tuple[str, str], asyncio.Task] = {}
        self._request_timeout = request_timeout
        self._reconnect_base = reconnect_base
        self._reconnect_max = reconnect_max

    async def startup(self) -> None:
        """Create the aiocoap client context."""
        self._context = await aiocoap.Context.create_client_context()
        logger.info("CoAP client context created")

    async def shutdown(self) -> None:
        """Cancel all subscriptions and close the context."""
        logger.info("Shutting down CoAP observer client...")

        # Cancel all subscription tasks
        for key, task in self._subscriptions.items():
            task.cancel()
            logger.debug("Cancelled subscription %s", key)
        self._subscriptions.clear()

        # Close context
        if self._context:
            await self._context.shutdown()
            self._context = None

        logger.info("CoAP observer client shut down")

    # ── Observe Subscriptions ──

    async def subscribe(
        self,
        device_id: str,
        ip: str,
        port: int,
        resource_uri: str,
    ) -> None:
        """
        Start an Observe subscription to a device resource.

        Runs as a background task with automatic reconnection on failure.
        """
        key = (device_id, resource_uri)
        if key in self._subscriptions and not self._subscriptions[key].done():
            logger.debug("Already subscribed to %s/%s", device_id, resource_uri)
            return

        task = asyncio.create_task(
            self._observe_loop(device_id, ip, port, resource_uri),
            name=f"observe-{device_id}-{resource_uri}",
        )
        self._subscriptions[key] = task
        logger.info("Started Observe subscription: %s/%s", device_id, resource_uri)

    async def unsubscribe(self, device_id: str, resource_uri: str) -> None:
        """Cancel a specific subscription."""
        key = (device_id, resource_uri)
        task = self._subscriptions.pop(key, None)
        if task and not task.done():
            task.cancel()
            logger.info("Unsubscribed from %s/%s", device_id, resource_uri)

    async def _observe_loop(
        self,
        device_id: str,
        ip: str,
        port: int,
        resource_uri: str,
    ) -> None:
        """
        Observe loop with exponential backoff on failure.

        Keeps reconnecting until cancelled.
        """
        backoff = self._reconnect_base

        while True:
            try:
                await self._run_observation(device_id, ip, port, resource_uri)
            except asyncio.CancelledError:
                logger.debug("Observe cancelled: %s/%s", device_id, resource_uri)
                return
            except Exception as e:
                logger.warning(
                    "Observe failed for %s/%s: %s — retrying in %.0fs",
                    device_id, resource_uri, e, backoff,
                )
                await self._store.mark_offline(device_id)
                if self._metrics:
                    self._metrics.record_coap_error(device_id, type(e).__name__)

                try:
                    await asyncio.sleep(backoff)
                except asyncio.CancelledError:
                    return

                backoff = min(backoff * 2, self._reconnect_max)
            else:
                # Observation ended normally (shouldn't happen) — reset backoff
                backoff = self._reconnect_base

    async def _run_observation(
        self,
        device_id: str,
        ip: str,
        port: int,
        resource_uri: str,
    ) -> None:
        """Execute a single Observe session."""
        uri = f"coap://{ip}:{port}/{resource_uri}"
        request = Message(code=Code.GET, uri=uri, observe=0)

        logger.debug("Sending Observe GET to %s", uri)

        try:
            response = await asyncio.wait_for(
                self._context.request(request).response,
                timeout=self._request_timeout,
            )
        except asyncio.TimeoutError:
            raise ConnectionError(f"Timeout connecting to {uri}")

        # Process the initial response
        await self._handle_notification(device_id, resource_uri, response)
        await self._store.mark_online(device_id)

        if self._metrics:
            self._metrics.record_coap_request(device_id)

        # Reset backoff on successful connection
        logger.info("Observe active: %s/%s", device_id, resource_uri)

        # Process subsequent notifications
        observation = self._context.request(Message(code=Code.GET, uri=uri, observe=0))

        async for notification in observation.observation:
            await self._handle_notification(device_id, resource_uri, notification)
            if self._metrics:
                self._metrics.record_coap_request(device_id)

    async def _handle_notification(
        self,
        device_id: str,
        resource_uri: str,
        response: Message,
    ) -> None:
        """Parse a CoAP response/notification and update the store."""
        try:
            payload = json.loads(response.payload.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            logger.warning(
                "Invalid payload from %s/%s: %s", device_id, resource_uri, e
            )
            return

        # Extract value and unit from the ESP JSON format
        # ESP sends: {"device": "...", "<resource_key>": value, "unit": "..."}
        value = None
        unit = payload.get("unit", "")

        # Map resource URIs to JSON keys
        resource_key_map = {
            "sensors/soil_moisture": "soil_moisture",
            "sensors/temperature": "temperature",
            "sensors/water_level": "water_level",
            "events/trigger": "trigger",
        }

        key = resource_key_map.get(resource_uri)
        if key and key in payload:
            value = payload[key]
        else:
            # Fallback: try "value" key or use full payload
            value = payload.get("value", payload)

        await self._store.update_reading(device_id, resource_uri, value, unit)

        if self._metrics:
            self._metrics.update_sensor_metric(device_id, resource_uri, value, unit)

        logger.debug(
            "Notification: %s/%s = %s %s", device_id, resource_uri, value, unit
        )

    # ── One-Shot Requests (for CoAP bridge) ──

    async def coap_get(
        self, ip: str, port: int, resource_uri: str
    ) -> dict[str, Any]:
        """
        Send a one-shot CoAP GET and return the parsed JSON payload.

        Raises ConnectionError on timeout, ValueError on invalid response.
        """
        uri = f"coap://{ip}:{port}/{resource_uri}"
        request = Message(code=Code.GET, uri=uri)

        try:
            response = await asyncio.wait_for(
                self._context.request(request).response,
                timeout=self._request_timeout,
            )
        except asyncio.TimeoutError:
            raise ConnectionError(f"CoAP GET timeout: {uri}")

        if not response.payload:
            return {"status": "ok", "code": str(response.code)}

        try:
            return json.loads(response.payload.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            raise ValueError(f"Invalid JSON from {uri}")

    async def coap_put(
        self, ip: str, port: int, resource_uri: str, payload: dict
    ) -> dict[str, Any]:
        """
        Send a one-shot CoAP PUT with a JSON payload.

        Raises ConnectionError on timeout, ValueError on invalid response.
        """
        uri = f"coap://{ip}:{port}/{resource_uri}"
        request = Message(
            code=Code.PUT,
            uri=uri,
            payload=json.dumps(payload).encode("utf-8"),
        )

        try:
            response = await asyncio.wait_for(
                self._context.request(request).response,
                timeout=self._request_timeout,
            )
        except asyncio.TimeoutError:
            raise ConnectionError(f"CoAP PUT timeout: {uri}")

        if not response.payload:
            return {"status": "ok", "code": str(response.code)}

        try:
            return json.loads(response.payload.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            raise ValueError(f"Invalid JSON from {uri}")

    # ── Status ──

    @property
    def active_subscriptions(self) -> int:
        """Number of currently active subscription tasks."""
        return sum(1 for t in self._subscriptions.values() if not t.done())

    def subscription_status(self) -> dict[str, str]:
        """Return status of each subscription (running/done/cancelled)."""
        result = {}
        for (dev_id, uri), task in self._subscriptions.items():
            key = f"{dev_id}/{uri}"
            if task.cancelled():
                result[key] = "cancelled"
            elif task.done():
                exc = task.exception()
                result[key] = f"failed: {exc}" if exc else "done"
            else:
                result[key] = "running"
        return result