"""
MeshBay Node — Hub client.

Handles all communication from the node to a Mesh Hub:
  - User registration (first run)
  - Login → JWT (access token + refresh token)
  - JWT offline verification and auto-refresh
  - Node announcement (endpoint_hint)
  - GEK bundle retrieval for a group
  - User public key lookup (for GEK wrapping)

JWT verification is done locally using the hub's cached Ed25519 public key.
The hub is only contacted for login and refresh — not for every request.
"""

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 cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives import serialization

from meshbay_common.crypto import pk_to_b64, unwrap_gek
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
    password: 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

    # ── Registration ──────────────────────────────────────────────────────────

    async def register(self) -> str:
        """Register this node's user on the hub. Returns user_id. Idempotent (409 ok)."""
        r = await self._http.post("/v1/users/register", json={
            "username":        self._config.username,
            "password":        self._config.password,
            "pk_user_ed25519": self._keys.pk_ed25519_b64,
            "pk_user_x25519":  self._keys.pk_x25519_b64,
        })
        if r.status_code == 201:
            log.info("Registered user '%s' on hub", self._config.username)
            return r.json()["user_id"]
        if r.status_code == 409:
            log.debug("User '%s' already registered", self._config.username)
            return ""
        r.raise_for_status()
        return ""

    # ── Login ─────────────────────────────────────────────────────────────────

    async def login(self) -> HubSession:
        """Login, verify JWT offline, return HubSession."""
        hub_pk_pem = await self._fetch_hub_pk()

        r = await self._http.post("/v1/users/login", json={
            "username": self._config.username,
            "password": self._config.password,
        })
        r.raise_for_status()
        data = r.json()

        access_token  = data["access_token"]
        refresh_token = data["refresh_token"]

        # Verify offline — if this passes, the hub's identity is confirmed
        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"

        self._session = HubSession(
            hub_url=self._config.hub_url,
            username=self._config.username,
            user_id=decoded["sub"],
            access_token=access_token,
            refresh_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 refresh_token(self) -> None:
        """Refresh the access token using the refresh token."""
        if self._session is None:
            raise RuntimeError("Not logged in")

        r = await self._http.post("/v1/users/token/refresh", json={
            "refresh_token": self._session.refresh_token,
        })
        r.raise_for_status()
        new_token = r.json()["access_token"]

        decoded = jwt.decode(new_token, self._session.hub_pk_pem, algorithms=["EdDSA"])
        self._session.access_token = new_token
        self._session._token_exp   = decoded["exp"]
        log.debug("Access token refreshed (exp in %ds)", self._session.token_expires_in)

    async def ensure_fresh_token(self) -> None:
        """Auto-refresh token if close to expiry."""
        if self._session and self._session.token_needs_refresh:
            await self.refresh_token()

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

    # ── GEK retrieval ─────────────────────────────────────────────────────────

    async def fetch_gek(self, group_id: str) -> bytes:
        """
        Fetch and unwrap the GEK bundle for a group.
        Returns the raw GEK bytes.
        """
        if self._session is None:
            raise RuntimeError("Not logged in")
        await self.ensure_fresh_token()

        r = await self._http.get(f"/v1/groups/{group_id}/gek",
                                 headers=self._session.auth_headers)
        if r.status_code == 404:
            raise LookupError(f"No GEK bundle found for group {group_id!r}")
        r.raise_for_status()

        bundle     = r.json()
        sk_x_raw   = self._keys.sk_x25519.private_bytes(
            serialization.Encoding.Raw, serialization.PrivateFormat.Raw,
            serialization.NoEncryption())
        pk_x_raw   = base64.b64decode(self._keys.pk_x25519_b64)

        gek = unwrap_gek(bundle, sk_x_raw, pk_x_raw)
        log.info("GEK unwrapped for group %s", group_id[:8])
        return gek

    # ── 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: register (idempotent) → login → announce node.
        Returns an active HubSession.
        """
        await self.register()
        session = await self.login()
        await self.announce_node(endpoint_hint)
        return session