"""Federation client for connecting to peer NEUROVAULT devices."""

from __future__ import annotations

import logging
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 FederationClient:
    """Connect to a peer NEUROVAULT device for bidirectional sync."""

    def __init__(self, device_id: str, key_pair, sync_engine):
        self._device_id = device_id
        self._key_pair = key_pair
        self._sync = sync_engine
        self._seq = 0

    async def sync_with_peer(
        self,
        host: str,
        port: int,
        since_checkpoint: str | None = None,
    ) -> dict[str, Any]:
        """Perform bidirectional sync with peer over WebSocket."""
        try:
            import websockets
        except Exception:  # pragma: no cover - optional dependency
            return {"error": "websockets not installed"}

        uri = f"ws://{host}:{int(port)}"

        try:
            async with websockets.connect(uri, max_size=50_000_000) as ws:
                raw_challenge = await ws.recv()
                challenge_msg = unpack_message(raw_challenge)
                challenge = str(
                    challenge_msg.get("payload", {}).get("challenge", "")
                    if isinstance(challenge_msg.get("payload"), dict)
                    else ""
                )
                if not challenge:
                    return {"error": "invalid challenge"}

                signature = self._key_pair.private_key.sign(challenge.encode("utf-8"))
                auth = {
                    "type": "AUTH_RESPONSE",
                    "from": self._device_id,
                    "payload": {
                        "public_key": self._key_pair.public_key_hex,
                        "signature": signature.hex(),
                    },
                }
                await ws.send(pack_message(auth))

                self._seq += 1
                request = {
                    "type": "SYNC_REQUEST",
                    "seq": self._seq,
                    "from": self._device_id,
                    "timestamp": int(time.time() * 1000),
                    "payload": {"since_checkpoint": since_checkpoint},
                }
                await ws.send(pack_message(request))

                raw_delta = await ws.recv()
                delta_msg = unpack_message(raw_delta)

                if delta_msg.get("type") != "SYNC_DELTA":
                    return {"error": "unexpected response", "type": delta_msg.get("type")}

                payload = delta_msg.get("payload", {}) if isinstance(delta_msg.get("payload"), dict) else {}
                bundle_hex = str(payload.get("bundle", ""))
                if not bundle_hex:
                    return {"error": "missing bundle"}

                bundle = KnowledgeBundle.from_msgpack(bytes.fromhex(bundle_hex))
                incoming_stats = self._sync.apply_bundle(bundle)

                our_bundle = self._sync.create_delta_bundle(
                    since_checkpoint=since_checkpoint,
                    device_id=self._device_id,
                )
                self._seq += 1
                our_delta = {
                    "type": "SYNC_DELTA",
                    "seq": self._seq,
                    "from": self._device_id,
                    "timestamp": int(time.time() * 1000),
                    "payload": {"bundle": our_bundle.to_msgpack().hex()},
                }
                await ws.send(pack_message(our_delta))

                ack_msg = unpack_message(await ws.recv())
                outgoing_ack = ack_msg.get("payload", {}) if isinstance(ack_msg.get("payload"), dict) else {}

                return {
                    "incoming": incoming_stats,
                    "outgoing_ack": outgoing_ack,
                    "peer": f"{host}:{int(port)}",
                }
        except Exception as exc:
            logger.warning("sync_with_peer failed: %s", exc)
            return {"error": "sync failed", "detail": str(exc), "peer": f"{host}:{int(port)}"}
