"""Home Assistant WebSocket client.

Implements the slice of the HA WebSocket API we need for Phase 1:

* Authenticate with a long-lived access token.
* Subscribe to events (optionally filtered by ``event_type``).
* Fetch the state list (``get_states``) for entity picker UI.
* Fetch the entity and area registries to build an ``entity_id -> area_id``
  lookup that the parser uses to enrich ``state_changed`` events with the
  area name.
* Run an indefinite subscription loop with exponential backoff reconnect.

The HA protocol reference is at
https://developers.home-assistant.io/docs/api/websocket/ — message ids are
ascending integers, server replies use the same id, and authentication must
complete before any other command is accepted.
"""

from __future__ import annotations

import asyncio
import itertools
import logging
import random
import time
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Awaitable, Callable
from urllib.parse import urlparse, urlunparse

import aiohttp

_LOGGER = logging.getLogger(__name__)


class HomeAssistantAuthError(Exception):
    """Raised when HA rejects our access token. Fatal — no point retrying."""


class HomeAssistantApiError(Exception):
    """Raised when an HA WS command returns ``success: false``."""


# Default reconnect backoff: 2s, 4s, 8s, ..., capped at 60s with jitter.
_RECONNECT_BASE_SECONDS = 2.0
_RECONNECT_MAX_SECONDS = 60.0
_RECONNECT_JITTER_RATIO = 0.2

# Bounded queue between the WS receive loop and the emit consumer. Overflow
# drops the oldest event (FIFO) and logs at WARNING — better to lose one
# state_changed than fall behind the firehose indefinitely.
_EMIT_QUEUE_SIZE = 1000


def _ws_url_from_base(base_url: str) -> str:
    """Derive the HA WebSocket URL from the user-provided HTTP(S) base URL.

    ``http://homeassistant.local:8123`` -> ``ws://homeassistant.local:8123/api/websocket``.
    The user enters their normal HA URL; we transform the scheme + append
    the API path. This keeps the UI single-field and avoids confusion about
    which URL form to use.

    Userinfo (``user:pass@host``) is **stripped** — credentials embedded in
    the URL would otherwise flow into log lines and exception strings via
    ``aiohttp`` error messages. The HA WS protocol uses an access-token
    handshake; HTTP basic auth in the URL is never the intended path.
    """
    parsed = urlparse(base_url.rstrip("/"))
    if parsed.scheme in ("ws", "wss"):
        scheme = parsed.scheme
    elif parsed.scheme == "https":
        scheme = "wss"
    else:
        scheme = "ws"
    # ``netloc`` may contain ``user:pass@host:port``; ``hostname`` + ``port``
    # rebuild it without the credential prefix.
    host = parsed.hostname or ""
    if parsed.port is not None:
        netloc = f"{host}:{parsed.port}"
    else:
        netloc = host
    return urlunparse(
        (scheme, netloc, "/api/websocket", "", "", "")
    )


def _redact(text: str) -> str:
    """Strip embedded credentials from text before logging.

    ``aiohttp`` exception strings include the URL, so a malformed
    ``https://token@host`` would otherwise expose the token. This is a
    defense-in-depth measure — ``_ws_url_from_base`` already strips
    userinfo from the connect URL, but third-party libs may quote the
    user-supplied input separately.
    """
    if not text:
        return text
    # Match ``scheme://[user[:pass]@]host`` and drop the userinfo segment.
    import re
    return re.sub(
        r"(?P<scheme>\w+://)(?:[^/@\s]+@)",
        r"\g<scheme>",
        text,
    )


class HomeAssistantWSClient:
    """Single-instance WebSocket client for one HA server."""

    def __init__(
        self,
        session: aiohttp.ClientSession,
        base_url: str,
        access_token: str,
        verify_tls: bool = True,
    ) -> None:
        self._session = session
        self._ws_url = _ws_url_from_base(base_url)
        self._access_token = access_token
        self._verify_tls = verify_tls
        self._id_counter = itertools.count(1)

    # ------------------------------------------------------------------
    # Connection primitives
    # ------------------------------------------------------------------

    @asynccontextmanager
    async def _connect(self) -> AsyncIterator[aiohttp.ClientWebSocketResponse]:
        """Open a fresh WS, complete the auth handshake, and yield the socket.

        Raises :class:`HomeAssistantAuthError` on invalid token (fatal) and
        :class:`HomeAssistantApiError` on other handshake failures (caller
        decides whether to retry).
        """
        ws = await self._session.ws_connect(
            self._ws_url,
            ssl=None if self._verify_tls else False,
            heartbeat=30,
            autoping=True,
        )
        try:
            await self._authenticate(ws)
            yield ws
        finally:
            await ws.close()

    async def _authenticate(self, ws: aiohttp.ClientWebSocketResponse) -> None:
        """Run the HA auth handshake on a freshly-opened socket."""
        greeting = await ws.receive_json(timeout=10)
        if greeting.get("type") != "auth_required":
            raise HomeAssistantApiError(
                f"Expected auth_required, got {greeting.get('type')!r}"
            )
        await ws.send_json({"type": "auth", "access_token": self._access_token})
        result = await ws.receive_json(timeout=10)
        msg_type = result.get("type")
        if msg_type == "auth_ok":
            return
        if msg_type == "auth_invalid":
            raise HomeAssistantAuthError(
                result.get("message") or "Home Assistant rejected the access token"
            )
        raise HomeAssistantApiError(
            f"Unexpected auth response: {msg_type!r}"
        )

    async def _send_command(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        payload: dict[str, Any],
    ) -> int:
        """Send a command with an auto-assigned id; return that id."""
        msg_id = next(self._id_counter)
        await ws.send_json({"id": msg_id, **payload})
        return msg_id

    async def _await_result(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        msg_id: int,
        timeout: float = 15.0,
    ) -> Any:
        """Wait for a ``result`` message matching ``msg_id`` and return its payload.

        ``time.monotonic`` is the right clock here — wall-clock deadlines
        would jump on NTP sync, and ``asyncio.get_event_loop().time()``
        is deprecated when called outside a running-loop context.
        """
        deadline = time.monotonic() + timeout
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise HomeAssistantApiError(
                    f"Timed out waiting for result of command id={msg_id}"
                )
            msg = await ws.receive_json(timeout=remaining)
            if msg.get("id") != msg_id:
                # Ignore unsolicited events that arrive between sending a
                # request-style command and its result.
                continue
            if msg.get("type") != "result":
                continue
            if not msg.get("success", False):
                err = msg.get("error", {})
                raise HomeAssistantApiError(
                    f"HA command failed: {err.get('code')} {err.get('message')}"
                )
            return msg.get("result")

    # ------------------------------------------------------------------
    # Multi-command session
    # ------------------------------------------------------------------

    @asynccontextmanager
    async def session(self) -> AsyncIterator["HomeAssistantSession"]:
        """Open one authenticated WS and let the caller run multiple commands.

        Each one-shot method (``get_states``, ``get_area_registry``, ...)
        opens a brand-new connection with a full TCP + WS + auth handshake.
        For callers that need to chain several queries (e.g. /status: connection
        check + entity list + area count) that overhead adds up — 3 separate
        TLS handshakes and 3 auth round-trips for what is really one logical
        request.

        Usage:

            async with client.session() as sess:
                states = await sess.get_states()
                areas = await sess.get_area_registry()

        The session shares the same id counter as the client, so message ids
        are unique across both one-shot calls and session-scoped calls if
        they happen to run concurrently against the same client instance.
        """
        async with self._connect() as ws:
            yield HomeAssistantSession(self, ws)

    # ------------------------------------------------------------------
    # One-shot commands
    # ------------------------------------------------------------------

    async def test_connection(self) -> tuple[bool, str]:
        """Connect, authenticate, and immediately close. Returns ``(ok, message)``."""
        try:
            async with self._connect() as _ws:
                return True, "OK"
        except HomeAssistantAuthError as err:
            return False, f"Auth failed: {err}"
        except (aiohttp.ClientError, asyncio.TimeoutError) as err:
            return False, f"Connection failed: {err}"
        except HomeAssistantApiError as err:
            return False, str(err)

    async def get_states(self) -> list[dict[str, Any]]:
        """Fetch the current state of every entity HA knows about."""
        async with self._connect() as ws:
            msg_id = await self._send_command(ws, {"type": "get_states"})
            result = await self._await_result(ws, msg_id)
            return list(result or [])

    async def get_area_registry(self) -> list[dict[str, Any]]:
        """Fetch the area registry (``area_id`` -> name + metadata)."""
        async with self._connect() as ws:
            msg_id = await self._send_command(
                ws, {"type": "config/area_registry/list"}
            )
            result = await self._await_result(ws, msg_id)
            return list(result or [])

    async def get_entity_registry(self) -> list[dict[str, Any]]:
        """Fetch the entity registry (entity_id -> area_id + metadata)."""
        async with self._connect() as ws:
            msg_id = await self._send_command(
                ws, {"type": "config/entity_registry/list"}
            )
            result = await self._await_result(ws, msg_id)
            return list(result or [])

    async def get_entity_to_area_lookup(self) -> dict[str, str]:
        """Build ``{entity_id: area_name}`` using the entity + area registries.

        Best-effort: returns an empty dict on any failure so the parser still
        works without area enrichment.
        """
        try:
            entities = await self.get_entity_registry()
            areas = await self.get_area_registry()
        except (HomeAssistantApiError, aiohttp.ClientError, asyncio.TimeoutError) as err:
            _LOGGER.warning("Could not fetch HA registry, areas disabled: %s", err)
            return {}
        area_names = {a.get("area_id"): a.get("name") for a in areas if a.get("area_id")}
        lookup: dict[str, str] = {}
        for entry in entities:
            entity_id = entry.get("entity_id")
            area_id = entry.get("area_id")
            if not isinstance(entity_id, str) or not area_id:
                continue
            name = area_names.get(area_id)
            if name:
                lookup[entity_id] = str(name)
        return lookup

    # ------------------------------------------------------------------
    # Subscription loop with reconnect
    # ------------------------------------------------------------------

    async def run_subscription(
        self,
        on_event: Callable[[dict[str, Any]], Awaitable[None]],
        event_types: list[str] | None = None,
        on_status_change: Callable[[str, str | None], None] | None = None,
        refresh_areas: Callable[[], Awaitable[dict[str, str]]] | None = None,
    ) -> None:
        """Run an indefinite subscription loop, reconnecting on drop.

        Parameters
        ----------
        on_event:
            Coroutine called with the inner ``event`` dict (the WS envelope is
            stripped). Slow callbacks apply TCP backpressure naturally; the
            internal queue prevents unbounded memory growth if the callback
            stalls.
        event_types:
            Restrict the subscription to these HA event types. ``None`` or
            empty subscribes to everything (very loud — only use for debug).
        on_status_change:
            Callback invoked with ``("connected", None)`` after a successful
            handshake and ``("disconnected", reason)`` when a connection drops.
            Useful for surfacing connection state in the event log.
        refresh_areas:
            Optional coroutine called on each (re)connect to refresh the
            area lookup. The result is not used by ``run_subscription``
            itself — the caller stores it where its ``on_event`` can read.
        """
        attempt = 0
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=_EMIT_QUEUE_SIZE)
        overflow_count = 0

        async def _drain() -> None:
            while True:
                evt = await queue.get()
                try:
                    await on_event(evt)
                except Exception:  # noqa: BLE001
                    _LOGGER.exception("on_event callback raised; continuing")
                finally:
                    queue.task_done()

        drain_task = asyncio.create_task(_drain(), name="ha-emit-drain")
        try:
            while True:
                try:
                    async with self._connect() as ws:
                        attempt = 0
                        if on_status_change is not None:
                            on_status_change("connected", None)
                        if refresh_areas is not None:
                            try:
                                # Note: refresh_areas opens its own WS in our
                                # current design (each one-shot command does).
                                # Fine for v1 — a few hundred ms once per
                                # (re)connect.
                                await refresh_areas()
                            except Exception:  # noqa: BLE001
                                _LOGGER.exception("Area refresh failed; continuing without")

                        # Subscribe. Passing per-event-type subscriptions is
                        # cheaper than subscribing to everything and filtering
                        # in Python — HA does the filtering.
                        if event_types:
                            for evt_type in event_types:
                                sub_id = await self._send_command(
                                    ws,
                                    {"type": "subscribe_events", "event_type": evt_type},
                                )
                                await self._await_result(ws, sub_id)
                        else:
                            sub_id = await self._send_command(
                                ws, {"type": "subscribe_events"}
                            )
                            await self._await_result(ws, sub_id)

                        async for msg in ws:
                            if msg.type == aiohttp.WSMsgType.TEXT:
                                payload = msg.json()
                                if payload.get("type") != "event":
                                    continue
                                event_obj = payload.get("event")
                                if not isinstance(event_obj, dict):
                                    continue
                                try:
                                    queue.put_nowait(event_obj)
                                except asyncio.QueueFull:
                                    overflow_count += 1
                                    if overflow_count % 50 == 1:
                                        _LOGGER.warning(
                                            "HA event queue full, dropped %d events so far "
                                            "(consumer is slower than HA event rate)",
                                            overflow_count,
                                        )
                                    # Drop oldest, retry put. This keeps the
                                    # most recent state visible at the cost
                                    # of older transient changes.
                                    try:
                                        queue.get_nowait()
                                        queue.task_done()
                                    except asyncio.QueueEmpty:
                                        pass
                                    try:
                                        queue.put_nowait(event_obj)
                                    except asyncio.QueueFull:
                                        pass
                            elif msg.type in (
                                aiohttp.WSMsgType.CLOSED,
                                aiohttp.WSMsgType.CLOSING,
                                aiohttp.WSMsgType.ERROR,
                            ):
                                raise aiohttp.ClientConnectionError(
                                    f"WS closed: {msg.type.name}"
                                )
                            else:
                                # PING/PONG handled by aiohttp autoping=True;
                                # BINARY/CONTINUATION are not used by HA today.
                                # Log at debug so a future protocol change is
                                # visible without spamming production logs.
                                _LOGGER.debug(
                                    "Ignored WS message of type %s", msg.type.name,
                                )
                except HomeAssistantAuthError as err:
                    # Fatal — caller must fix the access token. Reraise so
                    # the provider can mark itself unhealthy.
                    if on_status_change is not None:
                        on_status_change("disconnected", _redact(f"auth: {err}"))
                    raise
                except asyncio.CancelledError:
                    if on_status_change is not None:
                        on_status_change("disconnected", "cancelled")
                    raise
                except Exception as err:  # noqa: BLE001
                    redacted = _redact(str(err))
                    if on_status_change is not None:
                        on_status_change("disconnected", redacted)
                    delay = min(
                        _RECONNECT_BASE_SECONDS * (2 ** attempt),
                        _RECONNECT_MAX_SECONDS,
                    )
                    delay *= 1 + random.uniform(-_RECONNECT_JITTER_RATIO, _RECONNECT_JITTER_RATIO)
                    _LOGGER.warning(
                        "HA WS connection lost (%s); reconnecting in %.1fs",
                        redacted, delay,
                    )
                    attempt = min(attempt + 1, 10)
                    await asyncio.sleep(delay)
        finally:
            drain_task.cancel()
            # Drain task may finish via CancelledError (normal) or via an
            # unhandled exception thrown by on_event. Either way is fine here
            # — we're tearing down. Split the two cases for clarity rather
            # than catching `Exception` + `CancelledError` in one clause.
            try:
                await drain_task
            except asyncio.CancelledError:
                pass
            except Exception:  # noqa: BLE001
                _LOGGER.exception("HA drain task raised during shutdown")


# ---------------------------------------------------------------------------
# Multi-command session
# ---------------------------------------------------------------------------


class HomeAssistantSession:
    """A multi-command HA WS session bound to a single authenticated socket.

    Created via :meth:`HomeAssistantWSClient.session`. Use when you need to
    issue several commands in a row — sharing the connection saves the TCP
    + WS + auth round trips for every command after the first.

    The session forwards id assignment to the parent client's monotonic
    counter so ids stay unique across all sessions sharing the same client.
    """

    def __init__(
        self,
        client: HomeAssistantWSClient,
        ws: aiohttp.ClientWebSocketResponse,
    ) -> None:
        self._client = client
        self._ws = ws

    async def send(self, payload: dict[str, Any], timeout: float = 15.0) -> Any:
        """Send one command and wait for its ``result`` envelope."""
        msg_id = await self._client._send_command(self._ws, payload)
        return await self._client._await_result(self._ws, msg_id, timeout=timeout)

    async def get_states(self) -> list[dict[str, Any]]:
        result = await self.send({"type": "get_states"})
        return list(result or [])

    async def get_area_registry(self) -> list[dict[str, Any]]:
        result = await self.send({"type": "config/area_registry/list"})
        return list(result or [])

    async def get_entity_registry(self) -> list[dict[str, Any]]:
        result = await self.send({"type": "config/entity_registry/list"})
        return list(result or [])
