from __future__ import annotations

import asyncio
import json
import random
import time
from datetime import datetime, timezone

import aiohttp
import structlog

from ..config import Settings
from ..event_emitter import EventEmitter
from ..models import HAEventType, Severity
from .base import Monitor

_log = structlog.get_logger().bind(monitor="websocket")


class _AuthError(Exception):
    """Raised when HA returns auth_invalid during the WS handshake."""


def _make_ws_url(ha_url: str) -> str:
    if ha_url.startswith("https://"):
        base = ha_url.replace("https://", "wss://", 1)
    else:
        base = ha_url.replace("http://", "ws://", 1)
    return base.rstrip("/") + "/api/websocket"


class WebSocketMonitor(Monitor):
    """Persistent WebSocket connection to HA for real-time liveness monitoring.

    Subscribes to state_changed events — any traffic proves HA is alive.
    The watchdog fires ha_websocket_dead when the connection is silent for
    longer than silence_threshold, or immediately on disconnect.
    ha_websocket_recovered is emitted when the connection is restored after
    a dead alert was sent (allows supervisor to clear active incidents).
    """

    def __init__(
        self,
        ha_url: str,
        token: str,
        settings: Settings,
        emitter: EventEmitter,
        session: aiohttp.ClientSession,
    ) -> None:
        self._ws_url = _make_ws_url(ha_url)
        self._token = token
        self._settings = settings
        self._emitter = emitter
        self._session = session

        self._state: str = "disconnected"
        self._last_event_monotonic: float = time.monotonic()
        # 0.0 means no ha_websocket_dead has been emitted yet (for this session)
        self._last_dead_alert_at: float = 0.0

        self._stopping = False
        self._msg_id = 0
        self._main_task: asyncio.Task | None = None
        self._watchdog_task: asyncio.Task | None = None

    # ------------------------------------------------------------------
    # Monitor ABC
    # ------------------------------------------------------------------

    async def start(self) -> None:
        if not self._settings.websocket_enabled:
            _log.info("ws_monitor_disabled")
            return
        self._stopping = False
        self._last_event_monotonic = time.monotonic()
        self._main_task = asyncio.create_task(
            self._connection_loop(), name="ws_connection_loop"
        )
        self._watchdog_task = asyncio.create_task(
            self._watchdog_loop(), name="ws_watchdog"
        )
        _log.info("ws_monitor_started", ws_url=self._ws_url)

    async def stop(self) -> None:
        self._stopping = True
        self._state = "stopped"
        tasks = [t for t in [self._main_task, self._watchdog_task] if t is not None]
        for t in tasks:
            t.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        self._main_task = None
        self._watchdog_task = None
        _log.info("ws_monitor_stopped")

    @property
    def is_healthy(self) -> bool:
        if not self._settings.websocket_enabled:
            return True  # disabled monitors are not unhealthy
        return self._state == "subscribed"

    # ------------------------------------------------------------------
    # Connection loop — reconnects with exponential back-off
    # ------------------------------------------------------------------

    async def _connection_loop(self) -> None:
        delay = float(self._settings.websocket_reconnect_initial_delay)
        while not self._stopping:
            self._state = "connecting"
            clean_close = False
            try:
                await self._connect_and_listen()
                clean_close = True
                delay = float(self._settings.websocket_reconnect_initial_delay)
            except asyncio.CancelledError:
                raise
            except _AuthError as exc:
                _log.error("ws_auth_failed", error=str(exc))
                # Auth failures won't self-heal on fast retry — jump to max delay
                delay = float(self._settings.websocket_reconnect_max_delay)
            except Exception as exc:
                _log.warning("ws_connect_error", error=str(exc))

            self._state = "disconnected"
            if not self._stopping:
                self._on_disconnected()

            if self._stopping:
                break

            if clean_close:
                wait = 1.0  # brief pause before reconnecting after a clean HA close
            else:
                jitter_range = delay * self._settings.websocket_reconnect_jitter
                wait = max(0.1, delay + random.uniform(-jitter_range, jitter_range))
                delay = min(delay * 2, float(self._settings.websocket_reconnect_max_delay))

            _log.debug("ws_reconnect_wait", seconds=round(wait, 2))
            await asyncio.sleep(wait)

    # ------------------------------------------------------------------
    # Connect, auth, subscribe, receive
    # ------------------------------------------------------------------

    async def _connect_and_listen(self) -> None:
        # Override the session-level timeout: WS must stay open indefinitely,
        # only the initial TCP connect should be bounded.
        ws_timeout = aiohttp.ClientTimeout(total=None, connect=10.0, sock_connect=10.0)
        async with self._session.ws_connect(
            self._ws_url,
            timeout=ws_timeout,
            heartbeat=30.0,
        ) as ws:
            self._state = "authenticating"

            # Receive auth_required
            try:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=10.0)
            except (asyncio.TimeoutError, TypeError, json.JSONDecodeError) as exc:
                raise ConnectionError(f"Failed to receive auth_required: {exc}") from exc

            if msg.get("type") != "auth_required":
                raise ConnectionError(
                    f"Unexpected initial message type: {msg.get('type')!r}"
                )

            await ws.send_json({"type": "auth", "access_token": self._token})

            # Receive auth_ok or auth_invalid
            try:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=10.0)
            except (asyncio.TimeoutError, TypeError, json.JSONDecodeError) as exc:
                raise ConnectionError(f"Failed to receive auth response: {exc}") from exc

            if msg.get("type") == "auth_invalid":
                raise _AuthError(msg.get("message", "auth_invalid"))
            if msg.get("type") != "auth_ok":
                raise ConnectionError(
                    f"Unexpected auth response type: {msg.get('type')!r}"
                )

            # Subscribe to state_changed events
            self._msg_id += 1
            await ws.send_json({
                "id": self._msg_id,
                "type": "subscribe_events",
                "event_type": "state_changed",
            })

            # Mark connected — capture prior dead state before resetting
            prev_dead_at = self._last_dead_alert_at
            self._state = "subscribed"
            self._last_event_monotonic = time.monotonic()

            # Emit recovery if this reconnect follows a dead alert
            if prev_dead_at > 0.0:
                self._last_dead_alert_at = 0.0
                self._emit_recovered()

            _log.info("ws_subscribed", ws_url=self._ws_url)

            # Receive loop — any TEXT message proves HA is alive
            async for raw in ws:
                if self._stopping:
                    break
                if raw.type == aiohttp.WSMsgType.TEXT:
                    self._last_event_monotonic = time.monotonic()
                elif raw.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE):
                    _log.warning("ws_closed_by_server", msg_type=raw.type.name)
                    break

    # ------------------------------------------------------------------
    # Watchdog loop — detects silence while the WS appears connected
    # ------------------------------------------------------------------

    async def _watchdog_loop(self) -> None:
        while not self._stopping:
            try:
                await asyncio.sleep(self._settings.websocket_watchdog_interval_seconds)
            except asyncio.CancelledError:
                raise

            if self._state != "subscribed":
                continue  # disconnects are handled by the connection loop

            now = time.monotonic()
            silent_secs = now - self._last_event_monotonic
            if silent_secs <= self._settings.websocket_silence_threshold_seconds:
                continue

            cooldown = self._settings.websocket_down_alert_repeat_minutes * 60
            if self._last_dead_alert_at == 0.0 or (now - self._last_dead_alert_at) >= cooldown:
                self._emitter.emit(
                    event_type=HAEventType.ha_websocket_dead.value,
                    severity=Severity.error.value,
                    service="homeassistant",
                    message=(
                        f"HA WebSocket silent for {silent_secs:.0f}s — no events received"
                    ),
                    payload=self._dead_payload(silent_secs),
                )
                self._last_dead_alert_at = now
                _log.warning("ws_silent_dead_emitted", silent_seconds=round(silent_secs))

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _on_disconnected(self) -> None:
        """Emit ha_websocket_dead on connection loss, respecting cooldown."""
        if self._stopping:
            return
        now = time.monotonic()
        cooldown = self._settings.websocket_down_alert_repeat_minutes * 60
        if self._last_dead_alert_at == 0.0 or (now - self._last_dead_alert_at) >= cooldown:
            silent_secs = now - self._last_event_monotonic
            self._emitter.emit(
                event_type=HAEventType.ha_websocket_dead.value,
                severity=Severity.error.value,
                service="homeassistant",
                message=f"HA WebSocket disconnected — silent for {silent_secs:.0f}s",
                payload=self._dead_payload(silent_secs),
            )
            self._last_dead_alert_at = now
            _log.warning("ws_dead_emitted", silent_seconds=round(silent_secs))

    def _emit_recovered(self) -> None:
        self._emitter.emit(
            event_type=HAEventType.ha_websocket_recovered.value,
            severity=Severity.info.value,
            service="homeassistant",
            message="HA WebSocket reconnected and receiving events",
            payload={"connection_state": "subscribed"},
        )
        _log.info("ws_recovered_emitted")

    def _dead_payload(self, silent_secs: float) -> dict:
        event_age = time.monotonic() - self._last_event_monotonic
        last_event_wall = time.time() - event_age
        return {
            "silent_seconds": round(silent_secs),
            "last_event_at": datetime.fromtimestamp(
                last_event_wall, tz=timezone.utc
            ).isoformat(),
            "connection_state": self._state,
        }
