"""
MeshBay Node — Hub client.

Handles all communication from the node to a Mesh Hub:
  - Ed25519 authentication (node-scoped JWT, no password material on node)
  - JWT offline verification and auto-refresh
  - Node announcement (endpoint_hint)
  - User public key lookup (for GEK wrapping)
  - Swarm hash registration

The node authenticates via Ed25519 challenge-response (/v1/nodes/auth).
No auth_key or password is ever stored on or transmitted from the node.
The hub issues a node-scoped JWT that cannot manage group membership.
"""

import base64
import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

import httpx
import jwt

from meshbay_node.keystore import NodeKeys

log = logging.getLogger(__name__)

TOKEN_REFRESH_MARGIN = 300   # refresh access token 5 min before expiry


@dataclass
class HubSession:
    hub_url:       str
    username:      str
    user_id:       str
    access_token:  str
    refresh_token: str
    hub_pk_pem:    bytes         # cached hub Ed25519 public key
    node_id:       str = ""
    _token_exp:    int = 0

    @property
    def auth_headers(self) -> dict:
        return {"Authorization": f"Bearer {self.access_token}"}

    @property
    def token_expires_in(self) -> int:
        return max(0, self._token_exp - int(time.time()))

    @property
    def token_needs_refresh(self) -> bool:
        return self.token_expires_in < TOKEN_REFRESH_MARGIN


@dataclass
class HubConfig:
    hub_url:  str
    username: str
    cache_dir: Path = field(default_factory=lambda: Path.home() / ".config" / "meshbay")

    @property
    def hub_pk_cache_path(self) -> Path:
        safe = self.hub_url.replace("://", "_").replace("/", "_").replace(":", "_")
        return self.cache_dir / f"hub_pk_{safe}.pem"


# ── Hub client ────────────────────────────────────────────────────────────────

class HubClient:
    """Async hub client. Use as async context manager or call close() explicitly."""

    def __init__(self, config: HubConfig, keys: NodeKeys):
        self._config  = config
        self._keys    = keys
        self._http    = httpx.AsyncClient(timeout=15, base_url=config.hub_url)
        self._session: HubSession | None = None
        self._ws: Any = None

    async def __aenter__(self):
        return self

    async def __aexit__(self, *_):
        await self.close()

    async def close(self):
        await self._http.aclose()

    # ── Hub public key ────────────────────────────────────────────────────────

    async def _fetch_hub_pk(self) -> bytes:
        """Fetch and cache hub Ed25519 public key PEM."""
        cache = self._config.hub_pk_cache_path
        if cache.exists():
            log.debug("Hub PK loaded from cache: %s", cache)
            return cache.read_bytes()

        r = await self._http.get("/v1/hub/pubkey")
        r.raise_for_status()
        pem = r.json()["pk_hub_pem"].encode()

        cache.parent.mkdir(parents=True, exist_ok=True)
        cache.write_bytes(pem)
        log.info("Hub PK fetched and cached: %s", cache)
        return pem

    # ── Ed25519 authentication ───────────────────────────────────────────────

    async def login(self) -> HubSession:
        """Authenticate via Ed25519 challenge-response. Returns node-scoped HubSession."""
        hub_pk_pem = await self._fetch_hub_pk()

        timestamp = int(time.time())
        message = f"meshbay:node_auth:{self._config.username}:{timestamp}".encode()
        signature = self._keys.sk_ed25519.sign(message)

        r = await self._http.post("/v1/nodes/auth", json={
            "username":  self._config.username,
            "timestamp": timestamp,
            "signature": base64.b64encode(signature).decode(),
        })
        r.raise_for_status()
        data = r.json()

        access_token = data["access_token"]

        decoded = jwt.decode(access_token, hub_pk_pem, algorithms=["EdDSA"])
        assert decoded["pk_user"] == self._keys.pk_ed25519_b64, \
            "Hub returned token for wrong public key"
        assert "jti" in decoded, "Hub token missing jti — hub is outdated"
        assert decoded.get("scope") == "node", \
            "Expected node-scoped token"

        if self._session:
            self._session.access_token = access_token
            self._session._token_exp = decoded["exp"]
        else:
            self._session = HubSession(
                hub_url=self._config.hub_url,
                username=self._config.username,
                user_id=decoded["sub"],
                access_token=access_token,
                refresh_token="",
                hub_pk_pem=hub_pk_pem,
                _token_exp=decoded["exp"],
            )
        log.info("Logged in as '%s' (exp in %ds)", self._config.username,
                 self._session.token_expires_in)
        return self._session

    async def ensure_fresh_token(self) -> None:
        """Re-authenticate with Ed25519 if token is close to expiry."""
        if self._session and self._session.token_needs_refresh:
            await self.login()

    # ── Node announcement ─────────────────────────────────────────────────────

    async def announce_node(self, endpoint_hint: str | None = None) -> str:
        """Announce this node to the hub. Returns node_id."""
        if self._session is None:
            raise RuntimeError("Not logged in")
        await self.ensure_fresh_token()

        r = await self._http.post("/v1/nodes/announce", json={
            "pk_node":       self._keys.pk_ed25519_b64,
            "endpoint_hint": endpoint_hint,
        }, headers=self._session.auth_headers)
        r.raise_for_status()
        node_id = r.json()["node_id"]
        self._session.node_id = node_id
        log.info("Node announced: %s (hint=%s)", node_id[:8], endpoint_hint)
        return node_id

    # ── User pubkey lookup ────────────────────────────────────────────────────

    async def get_user_pubkeys(self, username: str) -> dict:
        """Return {'pk_ed25519': str, 'pk_x25519': str} for a user."""
        if self._session is None:
            raise RuntimeError("Not logged in")
        await self.ensure_fresh_token()

        r = await self._http.get(f"/v1/users/{username}/pubkeys",
                                 headers=self._session.auth_headers)
        if r.status_code == 404:
            raise LookupError(f"User not found: {username!r}")
        r.raise_for_status()
        return r.json()

    # ── Persistent WebSocket (signaling + revocations) ──────────────────────

    async def send_ws(self, data: str) -> None:
        """Send a message on the hub WebSocket (if connected). Best-effort."""
        ws = self._ws
        if ws:
            try:
                await ws.send(data)
            except Exception:
                pass

    async def maintain_ws(
        self,
        on_incoming: Any = None,
        on_revocation: Any = None,
        on_webrtc_offer: Any = None,
        group_ids: list[str] | None = None,
    ) -> None:
        """
        Maintain a persistent WebSocket connection to the hub.
        Receives NAT punch requests, revocation tokens, and WebRTC offers.
        Runs until cancelled.
        """
        import websockets

        if self._session is None:
            raise RuntimeError("Not logged in")

        hub_url = self._session.hub_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{hub_url}/v1/nodes/ws"

        while True:
            try:
                async with websockets.connect(ws_url) as ws:
                    auth_msg = {
                        "type": "auth",
                        "token": self._session.access_token,
                        "node_id": self._session.node_id,
                    }
                    if group_ids:
                        auth_msg["group_ids"] = group_ids
                    await ws.send(json.dumps(auth_msg))
                    auth_resp = json.loads(await ws.recv())
                    if auth_resp.get("type") != "auth_ok":
                        log.error("WS auth failed: %s", auth_resp)
                        return

                    self._ws = ws
                    log.info("Hub WS connected")

                    async for raw in ws:
                        msg = json.loads(raw)
                        mtype = msg.get("type")

                        if mtype == "client_incoming" and on_incoming:
                            await on_incoming(msg["peer_ip"], msg["peer_port"])
                            await ws.send(json.dumps({"type": "punch_ready"}))

                        elif mtype == "revocation" and on_revocation:
                            on_revocation(msg.get("token", ""))

                        elif mtype == "webrtc_offer" and on_webrtc_offer:
                            answer = await on_webrtc_offer(
                                msg["sdp"], msg["peer_id"],
                                msg.get("ice_candidates", []))
                            if answer:
                                await ws.send(json.dumps({
                                    "type": "webrtc_answer",
                                    "peer_id": msg["peer_id"],
                                    "sdp": answer[0],
                                    "ice_candidates": answer[1],
                                }))

                        elif mtype == "pong":
                            pass

            except asyncio.CancelledError:
                raise
            except Exception as e:
                log.warning("Hub WS disconnected: %s — reconnecting in 5s", e)
                await asyncio.sleep(5)
            finally:
                self._ws = None

    # ── Swarm registration ─────────────────────────────────────────────────

    async def register_swarm(self, content_hashes: list[str], endpoint: str) -> int:
        """Register file hashes in the hub swarm table. Returns count registered."""
        if self._session is None:
            raise RuntimeError("Not logged in")
        await self.ensure_fresh_token()

        registered = 0
        for h in content_hashes:
            try:
                r = await self._http.post("/v1/swarm/register", json={
                    "content_hash": h,
                    "endpoint": endpoint,
                }, headers=self._session.auth_headers)
                if r.status_code in (201, 200):
                    registered += 1
            except Exception:
                pass
        return registered

    # ── Convenience: full startup sequence ───────────────────────────────────

    async def startup(self, endpoint_hint: str | None = None) -> HubSession:
        """
        Full startup sequence: Ed25519 login → announce node.
        The operator must register separately (browser or setup script).
        """
        session = await self.login()
        await self.announce_node(endpoint_hint)
        return session