"""Async BridgeClient — Ed25519 authenticated WebSocket connection to the Bridge.

Mirrors the Node.js ``bridge-client-node/src/bridge-client.ts`` implementation:
  - Connects to ``{bridge_url}/agent``
  - Sends ``agent_auth`` with Ed25519 signature, waits for ``auth_ok``
  - Automatic reconnection with exponential back-off + jitter
  - WebSocket ping/pong heartbeat
  - Emits events via callback hooks
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import math
import random
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Awaitable, Callable

import websockets
import websockets.asyncio.client
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed

from .keys import KeyPair, fingerprint, load_key_pair, sign

logger = logging.getLogger("beeos_bridge_client")


class ClientState(str, Enum):
    IDLE = "idle"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    CLOSED = "closed"


@dataclass
class ReconnectConfig:
    enabled: bool = True
    max_retries: int | float = math.inf
    base_delay: float = 1.0    # seconds
    max_delay: float = 30.0    # seconds


@dataclass
class HeartbeatConfig:
    interval: float = 30.0     # seconds
    timeout: float = 10.0      # seconds


EventCallback = Callable[..., Any]


@dataclass
class BridgeClientOptions:
    bridge_url: str
    service: str
    key_file: str | None = None
    private_key: bytes | None = None
    public_key: bytes | None = None
    reconnect: ReconnectConfig = field(default_factory=ReconnectConfig)
    heartbeat: HeartbeatConfig = field(default_factory=HeartbeatConfig)


class BridgeClient:
    """Async Ed25519-authenticated WebSocket client for the BeeOS Bridge."""

    def __init__(self, options: BridgeClientOptions) -> None:
        self._bridge_url = options.bridge_url.rstrip("/")
        self._service = options.service
        self._reconnect = options.reconnect
        self._heartbeat = options.heartbeat

        if options.private_key and options.public_key:
            self._keys = KeyPair(
                public_key=options.public_key,
                private_key=options.private_key,
            )
        elif options.key_file:
            self._keys = load_key_pair(options.key_file)
        else:
            raise ValueError(
                "BridgeClient requires either key_file or private_key+public_key"
            )

        self._ws: ClientConnection | None = None
        self._state = ClientState.IDLE
        self._retry_count = 0
        self._closing = False

        self._reconnect_task: asyncio.Task[None] | None = None
        self._heartbeat_task: asyncio.Task[None] | None = None
        self._receive_task: asyncio.Task[None] | None = None

        # Event callbacks: on_connected, on_disconnected, on_message, on_error
        self._listeners: dict[str, list[EventCallback]] = {}

    # -- Public properties -----------------------------------------------------

    @property
    def state(self) -> ClientState:
        return self._state

    @property
    def public_key_b64(self) -> str:
        """Base64 encoded public key — the agent's sole identity."""
        return base64.b64encode(self._keys.public_key).decode()

    @property
    def public_key_fingerprint(self) -> str:
        """SHA-256 hex fingerprint of the public key."""
        return fingerprint(self._keys.public_key)

    # -- Event system ----------------------------------------------------------

    def on(self, event: str, callback: EventCallback) -> BridgeClient:
        """Register an event listener. Events: connected, disconnected, message, error."""
        self._listeners.setdefault(event, []).append(callback)
        return self

    def _emit(self, event: str, *args: Any) -> None:
        for cb in self._listeners.get(event, []):
            try:
                result = cb(*args)
                if asyncio.iscoroutine(result):
                    asyncio.ensure_future(result)
            except Exception:
                logger.exception("Error in %s listener", event)

    # -- Connection lifecycle --------------------------------------------------

    async def connect(self) -> None:
        """Connect (or reconnect) to the Bridge."""
        if self._state in (ClientState.CONNECTED, ClientState.CONNECTING):
            return
        self._closing = False
        await self._do_connect()

    async def close(self, graceful: bool = True) -> None:
        """Close the connection and stop reconnection."""
        self._closing = True
        self._state = ClientState.CLOSED

        if self._reconnect_task and not self._reconnect_task.done():
            self._reconnect_task.cancel()
            self._reconnect_task = None

        self._stop_heartbeat()
        self._stop_receive()

        if self._ws:
            ws = self._ws
            self._ws = None
            try:
                if graceful:
                    await ws.close(1000, "client closing")
                else:
                    ws.transport.close()  # type: ignore[union-attr]
            except Exception:
                pass

    async def send(self, data: str | bytes) -> None:
        """Send a message to the Bridge. Raises if not connected."""
        if self._state != ClientState.CONNECTED or not self._ws:
            raise RuntimeError("Not connected")
        await self._ws.send(data)

    # -- Internal: connect + auth ----------------------------------------------

    async def _do_connect(self) -> None:
        self._state = (
            ClientState.CONNECTING if self._retry_count == 0 else ClientState.RECONNECTING
        )

        url = f"{self._bridge_url}/agent"
        try:
            ws = await websockets.asyncio.client.connect(
                url,
                ping_interval=None,  # we manage heartbeat ourselves
                close_timeout=5,
            )
        except Exception as exc:
            logger.warning("WebSocket connect failed: %s", exc)
            self._emit("error", exc)
            if self._state == ClientState.CONNECTING:
                raise
            self._schedule_reconnect()
            return

        self._ws = ws

        try:
            await self._send_auth(ws)
        except Exception as exc:
            logger.warning("Auth handshake failed: %s", exc)
            await ws.close()
            self._ws = None
            self._emit("error", exc)
            if self._state == ClientState.CONNECTING:
                raise
            self._schedule_reconnect()
            return

        self._state = ClientState.CONNECTED
        self._retry_count = 0
        self._start_heartbeat(ws)
        self._start_receive(ws)
        self._emit("connected")
        logger.info("Connected to Bridge (%s)", url)

    async def _send_auth(self, ws: ClientConnection) -> None:
        timestamp = int(time.time())
        nonce = str(uuid.uuid4())
        pub_b64 = self.public_key_b64

        msg = f"{self._service}|{pub_b64}|{timestamp}|{nonce}"
        signature = sign(msg, self._keys.private_key)

        auth_payload = json.dumps({
            "type": "agent_auth",
            "service": self._service,
            "public_key": pub_b64,
            "timestamp": timestamp,
            "nonce": nonce,
            "signature": base64.b64encode(signature).decode(),
        })

        await ws.send(auth_payload)

        # Wait for auth_ok (with 10s timeout)
        try:
            resp_raw = await asyncio.wait_for(ws.recv(), timeout=10.0)
        except asyncio.TimeoutError:
            raise RuntimeError("Auth timeout — no response within 10s")

        resp = json.loads(resp_raw)
        if resp.get("type") != "auth_ok":
            raise RuntimeError(f"Auth rejected: {resp}")

    # -- Internal: receive loop ------------------------------------------------

    def _start_receive(self, ws: ClientConnection) -> None:
        self._stop_receive()
        self._receive_task = asyncio.ensure_future(self._receive_loop(ws))

    def _stop_receive(self) -> None:
        if self._receive_task and not self._receive_task.done():
            self._receive_task.cancel()
            self._receive_task = None

    async def _receive_loop(self, ws: ClientConnection) -> None:
        try:
            async for raw in ws:
                try:
                    msg = json.loads(raw)
                    if msg.get("type") == "auth_ok":
                        continue
                except (json.JSONDecodeError, TypeError):
                    pass
                self._emit("message", raw)
        except ConnectionClosed as exc:
            reason = str(exc)
            logger.info("Disconnected: %s", reason)
        except asyncio.CancelledError:
            return
        except Exception as exc:
            self._emit("error", exc)
            reason = str(exc)
        else:
            reason = "connection closed"

        self._stop_heartbeat()
        self._state = ClientState.IDLE
        self._emit("disconnected", reason)

        if not self._closing and self._reconnect.enabled:
            self._schedule_reconnect()

    # -- Internal: heartbeat ---------------------------------------------------

    def _start_heartbeat(self, ws: ClientConnection) -> None:
        self._stop_heartbeat()
        self._heartbeat_task = asyncio.ensure_future(self._heartbeat_loop(ws))

    def _stop_heartbeat(self) -> None:
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            self._heartbeat_task = None

    async def _heartbeat_loop(self, ws: ClientConnection) -> None:
        """Send WebSocket pings and enforce pong timeout."""
        try:
            while True:
                await asyncio.sleep(self._heartbeat.interval)
                try:
                    pong = await ws.ping()
                    await asyncio.wait_for(pong, timeout=self._heartbeat.timeout)
                except (asyncio.TimeoutError, ConnectionClosed):
                    logger.warning("Heartbeat timeout, terminating connection")
                    try:
                        ws.transport.close()  # type: ignore[union-attr]
                    except Exception:
                        pass
                    return
        except asyncio.CancelledError:
            return

    # -- Internal: reconnect ---------------------------------------------------

    def _schedule_reconnect(self) -> None:
        if self._retry_count >= self._reconnect.max_retries:
            self._state = ClientState.CLOSED
            self._emit("error", RuntimeError("Max reconnect retries exceeded"))
            return

        jitter = random.uniform(0.85, 1.15)
        delay = min(
            self._reconnect.base_delay * (2 ** self._retry_count) * jitter,
            self._reconnect.max_delay,
        )
        self._retry_count += 1
        logger.info("Reconnecting in %.1fs (attempt %d)", delay, self._retry_count)

        async def _reconnect_after_delay() -> None:
            await asyncio.sleep(delay)
            try:
                await self._do_connect()
            except Exception as exc:
                self._emit("error", exc)

        self._reconnect_task = asyncio.ensure_future(_reconnect_after_delay())
