"""
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 asyncio
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()

        # Proof of possession of the node key (M8) — same domain-separated shape
        # as node_auth, so a signature for one can never satisfy the other.
        timestamp = int(time.time())
        message = (f"meshbay:node_announce:{self._session.user_id}:"
                   f"{self._keys.pk_ed25519_b64}:{timestamp}").encode()
        signature = base64.b64encode(
            self._keys.sk_ed25519.sign(message)).decode()

        r = await self._http.post("/v1/nodes/announce", json={
            "pk_node":       self._keys.pk_ed25519_b64,
            "endpoint_hint": endpoint_hint,
            "timestamp":     timestamp,
            "signature":     signature,
        }, 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"

        # Offers are handled off the read loop (see below), so keep a handle on
        # the tasks to avoid them being garbage-collected mid-negotiation.
        pending: set[asyncio.Task] = set()

        while True:
            try:
                # Explicit keepalive: this connection is how a node stays visible
                # to the hub, and a silently half-open socket looks exactly like a
                # working one until someone notices the node has vanished.
                async with websockets.connect(
                    ws_url, ping_interval=20, ping_timeout=20, close_timeout=5,
                    open_timeout=15,
                ) 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))
                    # Bounded: a hub that accepts the socket and then says nothing
                    # — which is what it does for a few seconds while restarting —
                    # would otherwise park this task here forever, with the node
                    # running, silent, and invisible to everyone.
                    auth_resp = json.loads(
                        await asyncio.wait_for(ws.recv(), timeout=15))
                    if auth_resp.get("type") != "auth_ok":
                        # Not fatal: the token may simply have expired while we
                        # were disconnected. Refresh on the next pass rather than
                        # ending the task, which used to strand the node for good.
                        log.warning("WS auth refused: %s — retrying in 5s", auth_resp)
                        await asyncio.sleep(5)
                        await self.ensure_fresh_token()
                        continue

                    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:
                            # Answered off the read loop on purpose. Awaiting the
                            # handler here meant one slow negotiation stopped the
                            # node reading this socket at all: no pings answered,
                            # no close frame noticed, no further offers served. A
                            # client that gave up mid-ICE left the node in
                            # CLOSE-WAIT, still running but invisible to the hub
                            # and unreachable by everyone, until it was restarted.
                            task = asyncio.create_task(
                                self._answer_offer(ws, on_webrtc_offer, msg))
                            pending.add(task)
                            task.add_done_callback(pending.discard)

                        elif mtype == "pong":
                            pass

            except asyncio.CancelledError:
                for task in pending:
                    task.cancel()
                raise
            except Exception as e:
                log.warning("Hub WS disconnected: %s — reconnecting in 5s", e)
                await asyncio.sleep(5)
            else:
                # A clean close ends the `async for` without raising. Say so, so a
                # node that quietly stopped being reachable leaves a trace.
                log.warning("Hub WS closed by the hub — reconnecting in 5s")
                await asyncio.sleep(5)
            finally:
                self._ws = None

    async def _answer_offer(self, ws, on_webrtc_offer, msg: dict) -> None:
        """Negotiate one WebRTC offer and return the answer, off the read loop."""
        try:
            answer = await on_webrtc_offer(
                msg["sdp"], msg["peer_id"], msg.get("ice_candidates", []))
        except Exception as e:
            log.warning("WebRTC offer from %s failed: %s",
                        str(msg.get("peer_id"))[:8], e)
            return
        if not answer:
            return
        try:
            await ws.send(json.dumps({
                "type": "webrtc_answer",
                "peer_id": msg["peer_id"],
                "sdp": answer[0],
                "ice_candidates": answer[1],
            }))
        except Exception as e:
            # The socket may have gone while we were negotiating; the client will
            # retry, and the read loop is reconnecting.
            log.warning("Could not deliver WebRTC answer to %s: %s",
                        str(msg.get("peer_id"))[:8], e)

    # ── 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