"""
MeshBay Node — revocation subscriber.

Maintains a persistent WebSocket connection to the hub.
When a signed revocation token arrives, verifies it (Ed25519)
and adds the revoked target to the local blocklist.

Usage in daemon:
    subscriber = RevocationSubscriber(hub_url, access_token, hub_pk_pem)
    await subscriber.start()          # connects in background
    # check membership:
    if subscriber.is_revoked("user", user_id):
        refuse connection
    await subscriber.stop()
"""

import asyncio
import json
import logging
import time
from typing import Literal

import httpx
import jwt

log = logging.getLogger(__name__)

RevocationTarget = Literal["user", "group"]


class RevocationSubscriber:
    """
    Background task that keeps a WebSocket connection to the hub
    and maintains a local revocation set.
    """

    def __init__(
        self,
        hub_url: str,
        access_token: str,
        hub_pk_pem: bytes,
        reconnect_delay: float = 5.0,
    ):
        self._hub_url        = hub_url.rstrip("/")
        self._access_token   = access_token
        self._hub_pk_pem     = hub_pk_pem
        self._reconnect_delay = reconnect_delay
        self._revoked_users:  set[str] = set()
        self._revoked_groups: set[str] = set()
        self._task: asyncio.Task | None = None
        self._running = False

    def is_revoked(self, target: RevocationTarget, target_id: str) -> bool:
        if target == "user":
            return target_id in self._revoked_users
        return target_id in self._revoked_groups

    def add_revocation(self, target: RevocationTarget, target_id: str) -> None:
        if target == "user":
            self._revoked_users.add(target_id)
            log.warning("User revoked locally: %s", target_id[:8])
        else:
            self._revoked_groups.add(target_id)
            log.warning("Group revoked locally: %s", target_id[:8])

    def verify_and_apply(self, token: str) -> bool:
        """Verify a revocation token and apply it. Returns True if valid."""
        try:
            payload = jwt.decode(token, self._hub_pk_pem, algorithms=["EdDSA"],
                                 options={"verify_exp": False})
            if payload.get("type") != "revocation":
                return False
            target    = payload["target"]
            target_id = payload["target_id"]
            self.add_revocation(target, target_id)
            return True
        except Exception as e:
            log.error("Invalid revocation token: %s", e)
            return False

    async def start(self) -> None:
        self._running = True
        self._task = asyncio.create_task(self._run_loop())
        log.info("RevocationSubscriber started")

    async def stop(self) -> None:
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        log.info("RevocationSubscriber stopped")

    async def _run_loop(self) -> None:
        while self._running:
            try:
                await self._connect_and_listen()
            except asyncio.CancelledError:
                raise
            except Exception as e:
                log.warning("WS disconnected (%s), reconnecting in %ss", e, self._reconnect_delay)
                await asyncio.sleep(self._reconnect_delay)

    async def _connect_and_listen(self) -> None:
        ws_url = self._hub_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url += "/v1/nodes/ws"

        async with httpx.AsyncClient() as client:
            async with client.stream("GET", ws_url,
                                     headers={"Upgrade": "websocket"}) as resp:
                # Use websockets library for proper WS protocol
                pass

        # Use websockets library directly
        import websockets
        async with websockets.connect(ws_url) as ws:
            # Authenticate
            await ws.send(json.dumps({
                "type":  "auth",
                "token": self._access_token,
            }))
            auth_resp = json.loads(await ws.recv())
            if auth_resp.get("type") != "auth_ok":
                raise ConnectionError(f"WS auth failed: {auth_resp}")
            log.info("WS connected to hub — node_id=%s", auth_resp.get("node_id", "?")[:8])

            # Listen for revocations + send keepalive pings
            ping_interval = 30.0
            last_ping = asyncio.get_event_loop().time()

            while self._running:
                now = asyncio.get_event_loop().time()
                if now - last_ping > ping_interval:
                    await ws.send(json.dumps({"type": "ping"}))
                    last_ping = now

                try:
                    msg_raw = await asyncio.wait_for(ws.recv(), timeout=ping_interval + 5)
                    msg = json.loads(msg_raw)

                    if msg.get("type") == "revocation":
                        token = msg.get("token", "")
                        ok = self.verify_and_apply(token)
                        log.info("Revocation received — valid=%s", ok)
                    elif msg.get("type") == "pong":
                        pass
                    else:
                        log.debug("WS message: %s", msg.get("type"))

                except asyncio.TimeoutError:
                    continue