"""Minimal in-process NIP-01 relay for tests.

Spins up ``websockets.serve`` on 127.0.0.1 on an OS-assigned free port. Each
instance owns its own asyncio loop in a daemon thread, so several relays can
run side-by-side in one test without sharing state.

Tests drive the relay via a handler callback ``(relay, ws, msg) -> None``
invoked per inbound NIP-01 frame. The handler may be sync or async and is
free to call ``ws.send(...)`` to push EVENT/EOSE/OK/CLOSED frames back. The
relay also records every received frame in ``relay.received`` for assertions.
"""

from __future__ import annotations

import asyncio
import json
import logging
import socket
import threading
from typing import Any, Callable

import websockets
from websockets.asyncio.server import serve as ws_serve


log = logging.getLogger("mock_relay")


class MockRelay:
    """Bound to 127.0.0.1 on an OS-assigned free port."""

    def __init__(
        self,
        handler: Callable[["MockRelay", Any, list], Any] | None = None,
    ):
        self.handler = handler or self._default_handler
        self.received: list[list] = []
        # Events the relay should replay to clients on every incoming REQ.
        # Populated by ``stage_event`` from test code; consumed by whichever
        # handler the test plugs in (the default fixture handler in
        # ``specs/conftest.py`` reads this list and sends ``EVENT`` frames
        # before its closing ``EOSE``).
        self.staged_events: list[dict] = []
        self.url: str | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._server: Any = None
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._stopped = threading.Event()
        self._connections: list[Any] = []

    @staticmethod
    def _free_port() -> int:
        with socket.socket() as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]

    def start(self) -> None:
        port = self._free_port()
        self.url = f"ws://127.0.0.1:{port}"
        log.debug("starting mock relay on %s", self.url)

        def _run():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            self._loop = loop

            async def _main():
                async with ws_serve(self._on_connect, "127.0.0.1", port) as srv:
                    self._server = srv
                    self._ready.set()
                    log.debug("mock relay %s ready", self.url)
                    while not self._stopped.is_set():
                        await asyncio.sleep(0.05)
                    log.debug("mock relay %s stop flag set", self.url)

            try:
                loop.run_until_complete(_main())
            finally:
                loop.close()
                log.debug("mock relay %s loop closed", self.url)

        self._thread = threading.Thread(
            target=_run, name="mock-relay", daemon=True,
        )
        self._thread.start()
        if not self._ready.wait(timeout=5.0):
            raise RuntimeError("mock relay didn't start")

    def stop(self) -> None:
        log.debug("stop() called for %s", self.url)
        self._stopped.set()
        if self._thread is not None:
            self._thread.join(timeout=3.0)

    async def _on_connect(self, ws) -> None:
        peer = getattr(ws, "remote_address", None)
        log.debug("mock relay %s: client connected from %s", self.url, peer)
        self._connections.append(ws)
        try:
            async for raw in ws:
                try:
                    msg = json.loads(raw)
                except Exception:
                    log.debug("mock relay %s: dropping non-json frame", self.url)
                    continue
                if not isinstance(msg, list):
                    log.debug("mock relay %s: dropping non-list frame: %r", self.url, msg)
                    continue
                self.received.append(msg)
                log.debug(
                    "mock relay %s: received frame type=%s body=%s",
                    self.url, msg[0] if msg else "?",
                    json.dumps(msg)[:200],
                )
                # Handler may be sync or async; if it returns a coroutine, await it.
                result = self.handler(self, ws, msg)
                if asyncio.iscoroutine(result):
                    await result
        except websockets.exceptions.ConnectionClosed as e:
            log.debug("mock relay %s: client closed (%s)", self.url, e)

    @staticmethod
    def _default_handler(_self: "MockRelay", _ws: Any, _msg: list) -> None:
        return None

    def stage_event(self, event: dict) -> None:
        """Append an event to the relay's replay buffer. The default handler
        used by spec fixtures replays every staged event on each incoming
        REQ — order matches insertion order. Tests stage events BEFORE the
        primary is subscribed so the very first REQ already sees them."""
        self.staged_events.append(event)

    def disconnect_all(self) -> None:
        """Forcibly drop every open client connection — simulates a network
        blip. Used by reconnect tests."""
        loop = self._loop
        assert loop is not None
        for ws in list(self._connections):
            asyncio.run_coroutine_threadsafe(
                ws.close(code=1011, reason="test-drop"), loop,
            )
