"""
MeshBay Hub — Revocation system.

Two components:
  1. WebSocket endpoint /v1/nodes/ws
     Nodes connect at startup and keep the connection alive.
     Hub pushes signed revocation tokens when a user or group is revoked.

  2. Admin endpoint POST /v1/admin/revoke
     Hub operator revokes a user or group.
     Signed revocation token is broadcast to all connected nodes.

Revocation token format (signed Ed25519):
  {
    "type":       "revocation",
    "target":     "user" | "group",
    "target_id":  "<user_id or group_id>",
    "reason":     "<reason string>",
    "revoked_at": <unix timestamp>,
    "jti":        "<uuid4>",
  }

Nodes verify the token with the hub's public key (already cached at startup).
On receipt: immediately refuse JWT tokens matching the revoked user_id,
and close active connections for that user.
"""

import base64
import json
import logging
import time
import uuid
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

import jwt
from meshbay_hub.auth import hub_public_key_pem, decode_access_token
from meshbay_hub.api.deps import get_current_user
from meshbay_hub.db.engine import get_db
from meshbay_hub.db.models import Group, IPLog, User

log = logging.getLogger(__name__)

router = APIRouter(tags=["revocation"])

# ── Connected node registry ───────────────────────────────────────────────────

_connected_nodes: dict[str, WebSocket] = {}   # node_id → websocket


def get_connected_node_count() -> int:
    return len(_connected_nodes)


async def broadcast_revocation(token: str) -> int:
    """Push a signed revocation token to all connected nodes. Returns count sent."""
    payload = json.dumps({"type": "revocation", "token": token})
    disconnected = []
    sent = 0
    for node_id, ws in _connected_nodes.items():
        try:
            await ws.send_text(payload)
            sent += 1
        except Exception:
            disconnected.append(node_id)
    for node_id in disconnected:
        _connected_nodes.pop(node_id, None)
    return sent


def _sign_revocation(target: str, target_id: str, reason: str) -> str:
    """Issue a signed revocation token (JWT EdDSA)."""
    from meshbay_hub.auth import _hub_sk_pem, _hub_id
    now = int(time.time())
    payload = {
        "type":       "revocation",
        "target":     target,       # "user" or "group"
        "target_id":  target_id,
        "reason":     reason,
        "revoked_at": now,
        "jti":        str(uuid.uuid4()),
        "iss":        _hub_id,
        "iat":        now,
    }
    return jwt.encode(payload, _hub_sk_pem, algorithm="EdDSA")


# ── WebSocket endpoint ────────────────────────────────────────────────────────

@router.websocket("/v1/nodes/ws")
async def node_websocket(ws: WebSocket):
    """
    Persistent WebSocket connection for nodes.
    Nodes authenticate with a JWT bearer in the first message.
    Hub sends revocation tokens as JSON messages.
    """
    await ws.accept()
    node_id: str | None = None

    try:
        # Auth: expect {"type": "auth", "token": "<jwt>"}
        raw = await ws.receive_text()
        msg = json.loads(raw)
        if msg.get("type") != "auth" or "token" not in msg:
            await ws.send_text(json.dumps({"type": "error", "detail": "Send auth first"}))
            await ws.close(code=4001)
            return

        try:
            decoded = decode_access_token(msg["token"])
        except Exception as e:
            await ws.send_text(json.dumps({"type": "error", "detail": str(e)}))
            await ws.close(code=4001)
            return

        node_id = decoded.get("sub", "unknown")
        _connected_nodes[node_id] = ws
        log.info("Node WS connected: %s", node_id[:8])
        await ws.send_text(json.dumps({"type": "auth_ok", "node_id": node_id}))

        # Keep-alive loop — wait for ping or disconnect
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            if msg.get("type") == "ping":
                await ws.send_text(json.dumps({"type": "pong"}))

    except WebSocketDisconnect:
        log.info("Node WS disconnected: %s", (node_id or "unknown")[:8])
    except Exception as e:
        log.error("Node WS error: %s", e)
    finally:
        if node_id:
            _connected_nodes.pop(node_id, None)


# ── Admin revocation endpoint ─────────────────────────────────────────────────

class RevokeRequest(BaseModel):
    target:    str   # "user" or "group"
    target_id: str
    reason:    str = "policy_violation"


@router.post("/v1/admin/revoke", status_code=200)
async def admin_revoke(
    body: RevokeRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """
    Revoke a user or group. Admin only (user must be hub admin — user_id in config).
    Issues a signed revocation token and broadcasts to all connected nodes.
    Also marks the target as revoked in the database.
    """
    if body.target not in ("user", "group"):
        raise HTTPException(status_code=422, detail="target must be 'user' or 'group'")

    # Mark as revoked in DB
    if body.target == "user":
        obj = await db.get(User, body.target_id)
        if not obj:
            raise HTTPException(status_code=404, detail="User not found")
        obj.status = "revoked"
    else:
        obj = await db.get(Group, body.target_id)
        if not obj:
            raise HTTPException(status_code=404, detail="Group not found")
        obj.status = "revoked"

    db.add(IPLog(
        user_id=current_user.id,
        event=f"revoke_{body.target}",
        ip_address="admin",
        detail=f"{body.target_id}: {body.reason}",
    ))
    await db.commit()

    # Issue and broadcast signed revocation token
    rev_token = _sign_revocation(body.target, body.target_id, body.reason)
    sent = await broadcast_revocation(rev_token)

    log.warning("Revoked %s %s — broadcast to %d nodes", body.target, body.target_id[:8], sent)
    return {
        "status":      "revoked",
        "target":      body.target,
        "target_id":   body.target_id,
        "nodes_notified": sent,
        "token":       rev_token,   # admin can store this for manual distribution
    }