"""WebSocket federation server for device-to-device sync."""

from __future__ import annotations

import asyncio
import logging
import os
import time
from typing import Any

from neurovault.network.bundle import KnowledgeBundle
from neurovault.network.codec import pack_message, unpack_message

logger = logging.getLogger(__name__)

class FederationServer:
    """WebSocket server implementing NEUROVAULT federation protocol."""

    def __init__(
        self,
        device_id: str,
        sync_engine,
        key_pair,
        trusted_peers: list[str] | None = None,
        port: int = 9473,
        host: str = "127.0.0.1",
    ):
        self._device_id = device_id
        self._sync = sync_engine
        self._key_pair = key_pair
        self._trusted = set(trusted_peers or [])
        self._port = int(port)
        self._host = host
        self._connections: dict[str, Any] = {}
        self._seq = 0
        self._server = None

    async def start(self) -> bool:
        try:
            import websockets
        except Exception:  # pragma: no cover - optional dependency
            logger.error("websockets package required for federation")
            return False

        self._server = await websockets.serve(self._handle_connection, self._host, self._port)
        logger.info("Federation server listening on %s:%d", self._host, self._port)
        return True

    async def run_forever(self) -> None:
        ok = await self.start()
        if not ok:
            return
        assert self._server is not None
        await self._server.wait_closed()

    async def stop(self) -> None:
        if self._server is None:
            return
        self._server.close()
        await self._server.wait_closed()
        self._server = None

    async def _handle_connection(self, ws, path=None):  # pragma: no cover - network I/O
        del path
        peer_id = None
        try:
            peer_id = await self._authenticate(ws)
            if not peer_id:
                await ws.close(code=4001, reason="Authentication failed")
                return

            self._connections[peer_id] = ws
            logger.info("Peer connected: %s", peer_id)

            async for raw_msg in ws:
                msg = unpack_message(raw_msg)
                await self._handle_message(ws, peer_id, msg)

        except Exception as exc:
            logger.warning("Connection error (peer %s): %s", peer_id, exc)
        finally:
            if peer_id:
                self._connections.pop(peer_id, None)

    async def _authenticate(self, ws) -> str | None:  # pragma: no cover - network I/O
        challenge = os.urandom(32).hex()

        hello_msg = self._make_message("CHALLENGE", {"challenge": challenge})
        await ws.send(pack_message(hello_msg))

        raw = await asyncio.wait_for(ws.recv(), timeout=10)
        response = unpack_message(raw)

        if response.get("type") != "AUTH_RESPONSE":
            return None

        peer_id = str(response.get("from", ""))
        payload = response.get("payload", {}) if isinstance(response.get("payload"), dict) else {}
        peer_pub_key_hex = str(payload.get("public_key", ""))
        signature_hex = str(payload.get("signature", ""))

        if self._trusted and peer_pub_key_hex not in self._trusted:
            logger.warning("Untrusted peer attempted connection: %s", peer_id)
            return None

        try:
            from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

            pub_key = Ed25519PublicKey.from_public_bytes(bytes.fromhex(peer_pub_key_hex))
            pub_key.verify(bytes.fromhex(signature_hex), challenge.encode("utf-8"))
            return peer_id
        except Exception as exc:
            logger.warning("Auth verification failed: %s", exc)
            return None

    async def _handle_message(self, ws, peer_id: str, msg: dict[str, Any]) -> None:  # pragma: no cover - network I/O
        msg_type = str(msg.get("type", ""))

        if msg_type == "SYNC_REQUEST":
            payload = msg.get("payload", {}) if isinstance(msg.get("payload"), dict) else {}
            since = payload.get("since_checkpoint")
            bundle = self._sync.create_delta_bundle(since_checkpoint=since, device_id=self._device_id)
            response = self._make_message("SYNC_DELTA", {"bundle": bundle.to_msgpack().hex()})
            await ws.send(pack_message(response))
            return

        if msg_type == "SYNC_DELTA":
            payload = msg.get("payload", {}) if isinstance(msg.get("payload"), dict) else {}
            bundle_hex = str(payload.get("bundle", ""))
            if not bundle_hex:
                return
            bundle = KnowledgeBundle.from_msgpack(bytes.fromhex(bundle_hex))
            stats = self._sync.apply_bundle(bundle)
            ack = self._make_message("SYNC_ACK", stats)
            await ws.send(pack_message(ack))
            return

        if msg_type == "HEARTBEAT":
            ack = self._make_message(
                "HEARTBEAT_ACK",
                {
                    "memory_count": self._sync._storage.count_memories(),
                    "uptime": int(time.time()),
                    "peer_id": peer_id,
                },
            )
            await ws.send(pack_message(ack))

    def _make_message(self, msg_type: str, payload: dict[str, Any]) -> dict[str, Any]:
        self._seq += 1
        return {
            "type": msg_type,
            "seq": self._seq,
            "from": self._device_id,
            "timestamp": int(time.time() * 1000),
            "payload": payload,
        }
