"""SwarmNode — minimal framework-agnostic wrapper over MeshNode.

Provides a thin, ergonomic wrapper around the lower-level ``MeshNode``
context manager. Exposes only the five operations that most framework
integrators need: connect, disconnect, publish, subscribe, discover, and
check_trust.

Example::

    from meshpocalypse_semantic import SwarmNode, SwarmNodeConfig

    config = SwarmNodeConfig(
        agent_id="my-agent",
        capabilities=["summarise", "search"],
    )
    node = SwarmNode(config)
    await node.connect()
    await node.publish("task.result", {"result": "done"})
    peers = await node.discover_capability("search")
    await node.disconnect()
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from typing import Any, Callable

import nats
from nats.aio.client import Client

from .atf import ATFLevel

logger = logging.getLogger(__name__)

# NATS subjects (must match the TypeScript SwarmNode constants)
_SUBJECT_AGENT_REGISTER = "mesh.agent.register"
_SUBJECT_CAPABILITIES_QUERY = "mesh.capabilities.query"
_SUBJECT_ATF_TRUST_CHECK = "mesh.atf.trust.check"

_REQUEST_TIMEOUT_S: float = 5.0


@dataclass
class SwarmNodeConfig:
    """Configuration for a :class:`SwarmNode`.

    Only ``agent_id`` and ``capabilities`` are required. All other fields
    have sensible defaults.

    Attributes:
        agent_id: Unique identifier for this agent.
        capabilities: Capabilities advertised by this agent (e.g.
            ``["summarise", "search"]``).
        trust_level: ATF trust level. Defaults to
            :attr:`~meshpocalypse_semantic.ATFLevel.SENIOR`.
        nats_url: NATS server URL. Defaults to ``"nats://localhost:4222"``.
        graph_url: Optional Oxigraph HTTP endpoint. When provided,
            :meth:`SwarmNode.connect` registers capabilities in the RDF
            graph via SPARQL UPDATE.
    """

    agent_id: str
    capabilities: list[str]
    trust_level: ATFLevel = ATFLevel.SENIOR
    nats_url: str = "nats://localhost:4222"
    graph_url: str | None = None


# Subscription handle returned by subscribe()
Unsubscribe = Callable[[], None]


class SwarmNode:
    """Minimal mesh participant.

    Connect, publish, discover — nothing else required.

    Wraps a NATS connection with an ergonomic API. Connects on
    :meth:`connect` and disconnects on :meth:`disconnect`. Can also be used
    as an async context manager.

    Args:
        config: :class:`SwarmNodeConfig` instance.

    Example::

        node = SwarmNode(SwarmNodeConfig(agent_id="a", capabilities=["x"]))
        await node.connect()
        try:
            await node.publish("greet", {"hello": "world"})
        finally:
            await node.disconnect()
    """

    def __init__(self, config: SwarmNodeConfig) -> None:
        self._config = config
        self._nc: Client | None = None
        self._subscriptions: list[Any] = []

    # ── Async context manager ────────────────────────────────────────────────

    async def __aenter__(self) -> "SwarmNode":
        """Connect on entry."""
        await self.connect()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        """Disconnect on exit."""
        await self.disconnect()

    # ── Public API ────────────────────────────────────────────────────────────

    async def connect(self) -> None:
        """Connect to NATS and register agent capabilities in the mesh.

        After calling ``connect()``, the node is reachable by other mesh
        participants. Capabilities are announced via a fire-and-forget
        publish to ``mesh.agent.register``.

        If ``config.graph_url`` is set, capabilities are also written to
        the Oxigraph RDF store. Graph write failures are logged but do not
        abort the connect flow.

        Raises:
            nats.errors.Error: If the NATS connection fails.
        """
        self._nc = await nats.connect(self._config.nats_url)

        payload = {
            "agentId": self._config.agent_id,
            "capabilities": self._config.capabilities,
            "trustLevel": self._config.trust_level.value,
            "timestamp": self._now_ms(),
        }
        await self._nc.publish(
            _SUBJECT_AGENT_REGISTER,
            json.dumps(payload).encode("utf-8"),
        )

        if self._config.graph_url:
            await self._register_capabilities_in_graph()

    async def disconnect(self) -> None:
        """Gracefully drain and disconnect from NATS.

        Cancels all active subscriptions before draining the connection.
        Safe to call multiple times.
        """
        if self._nc is None:
            return
        for sub in self._subscriptions:
            try:
                await sub.unsubscribe()
            except Exception:  # noqa: BLE001
                pass
        self._subscriptions.clear()
        try:
            await self._nc.drain()
        except Exception:  # noqa: BLE001
            pass
        self._nc = None

    async def publish(self, subject: str, data: dict[str, Any]) -> None:
        """Publish a JSON-serialisable payload to a NATS subject.

        Args:
            subject: NATS subject string (e.g. ``"task.result"``).
            data: JSON-serialisable dict to publish.

        Raises:
            RuntimeError: If the node is not connected.
            ValueError: If ``subject`` is empty.
        """
        self._require_connected()
        if not subject:
            raise ValueError("publish: subject must be a non-empty string")
        assert self._nc is not None  # narrowing for type checker
        await self._nc.publish(subject, json.dumps(data).encode("utf-8"))

    def subscribe(
        self,
        subject: str,
        handler: Callable[[dict[str, Any]], None],
    ) -> Unsubscribe:
        """Subscribe to a NATS subject.

        The ``handler`` is called with the decoded JSON payload for each
        message. Returns a zero-argument callable that cancels the
        subscription.

        Note:
            The subscription is registered asynchronously. Messages may be
            missed if they arrive before the first event-loop iteration after
            this call returns. For reliable operation, call this before
            :meth:`connect`, or await one event-loop tick after calling.

        Args:
            subject: NATS subject to subscribe to.
            handler: Callback invoked with the parsed JSON payload.

        Returns:
            An async-safe unsubscribe callable.

        Raises:
            RuntimeError: If the node is not connected.

        Example::

            def on_result(data: dict) -> None:
                print("got", data)

            unsubscribe = node.subscribe("task.result", on_result)
            # Later:
            unsubscribe()
        """
        self._require_connected()
        assert self._nc is not None

        import asyncio

        sub_holder: list[Any] = []

        async def _do_subscribe() -> None:
            async def _cb(msg: Any) -> None:
                try:
                    decoded = json.loads(msg.data.decode("utf-8"))
                except (json.JSONDecodeError, UnicodeDecodeError):
                    decoded = msg.data.decode("utf-8", errors="replace")
                handler(decoded)  # type: ignore[arg-type]

            sub = await self._nc.subscribe(subject, cb=_cb)  # type: ignore[union-attr]
            sub_holder.append(sub)
            self._subscriptions.append(sub)

        asyncio.ensure_future(_do_subscribe())

        def _unsubscribe() -> None:
            if sub_holder:
                sub = sub_holder[0]
                if sub in self._subscriptions:
                    self._subscriptions.remove(sub)
                asyncio.ensure_future(sub.unsubscribe())

        return _unsubscribe

    async def discover_capability(self, capability: str) -> list[str]:
        """Discover agents that advertise a given capability.

        Sends a NATS request to ``mesh.capabilities.query`` and returns the
        list of agent IDs that responded. Returns an empty list if no
        capability service is reachable within the timeout.

        Args:
            capability: Capability name to search for (e.g. ``"search"``).

        Returns:
            List of agent ID strings.
        """
        self._require_connected()
        assert self._nc is not None

        payload = json.dumps({"capability": capability}).encode("utf-8")
        try:
            response = await self._nc.request(
                _SUBJECT_CAPABILITIES_QUERY,
                payload,
                timeout=_REQUEST_TIMEOUT_S,
            )
            result: dict[str, Any] = json.loads(response.data.decode("utf-8"))
            agents = result.get("agents", [])
            return agents if isinstance(agents, list) else []
        except Exception as exc:
            logger.debug(
                "SwarmNode.discover_capability: NATS request failed — %s", exc
            )
            return []

    async def check_trust(self, peer_id: str) -> ATFLevel:
        """Check the ATF trust level of a peer agent.

        Sends a NATS request to ``mesh.atf.trust.check``. Returns
        :attr:`~meshpocalypse_semantic.ATFLevel.INTERN` (least privilege)
        if the ATF service is unreachable.

        Args:
            peer_id: Agent ID to check (e.g. ``"did:mesh:peer-agent"``).

        Returns:
            The peer's :class:`~meshpocalypse_semantic.ATFLevel`.
        """
        self._require_connected()
        assert self._nc is not None

        payload = json.dumps({"agentId": peer_id}).encode("utf-8")
        try:
            response = await self._nc.request(
                _SUBJECT_ATF_TRUST_CHECK,
                payload,
                timeout=_REQUEST_TIMEOUT_S,
            )
            result: dict[str, Any] = json.loads(response.data.decode("utf-8"))
            return ATFLevel(int(result["trustLevel"]))
        except Exception as exc:
            logger.debug("SwarmNode.check_trust: ATF request failed — %s", exc)
            return ATFLevel.INTERN

    # ── MCP server ───────────────────────────────────────────────────────────

    async def as_mcp_server(
        self,
        port: int = 3001,
        host: str = "0.0.0.0",
        tool_descriptions: dict[str, str] | None = None,
    ) -> None:
        """Expose this node capabilities as an MCP HTTP server.

        Starts an aiohttp-based JSON-RPC 2.0 server on host:port. Two endpoints
        are served:

        * POST /mcp  — MCP JSON-RPC 2.0 dispatcher (tools/list, tools/call).
        * GET  /mcp/health — Returns {"status": "ok", "toolCount": N}.

        The method blocks until the calling asyncio.Task is cancelled, at which
        point the server drains and shuts down cleanly. Wrap the call in a task
        if you need to run it concurrently::

            task = asyncio.create_task(node.as_mcp_server(port=3001))
            ...
            task.cancel()
            await task

        Args:
            port: TCP port to listen on. Defaults to 3001.
            host: Bind address. Defaults to 0.0.0.0 (all interfaces).
            tool_descriptions: Optional mapping of capability name to custom
                description string. Capabilities absent from this dict get an
                auto-generated description.

        Raises:
            OSError: If the port is already in use.
            asyncio.CancelledError: Re-raised after clean shutdown.
        """
        from .mcp_server import MCPServerConfig, start_mcp_server

        mcp_config = MCPServerConfig(
            agent_id=self._config.agent_id,
            capabilities=self._config.capabilities,
            nats_client=self._nc,
            tool_descriptions=dict(tool_descriptions or {}),
            host=host,
            port=port,
        )
        await start_mcp_server(mcp_config)

    # ── Properties ────────────────────────────────────────────────────────────

    @property
    def agent_id(self) -> str:
        """The configured agent ID."""
        return self._config.agent_id

    @property
    def capabilities(self) -> list[str]:
        """The capabilities advertised by this node."""
        return self._config.capabilities

    @property
    def trust_level(self) -> ATFLevel:
        """The configured ATF trust level."""
        return self._config.trust_level

    # ── Private helpers ───────────────────────────────────────────────────────

    def _require_connected(self) -> None:
        """Raise RuntimeError if the node is not connected."""
        if self._nc is None:
            raise RuntimeError(
                "SwarmNode not connected — call await node.connect() first, "
                "or use 'async with SwarmNode(config) as node:'"
            )

    @staticmethod
    def _now_ms() -> int:
        """Return current time as milliseconds since epoch."""
        import time

        return int(time.time() * 1000)

    async def _register_capabilities_in_graph(self) -> None:
        """Write capabilities to Oxigraph via SPARQL UPDATE over HTTP.

        Uses DELETE WHERE + INSERT DATA for idempotent registration.
        Failures are logged at WARNING level and do not abort the connect
        flow — the graph is an optional enhancement.
        """
        import urllib.request

        graph_url = (self._config.graph_url or "").rstrip("/")
        base_uri = f"mesh:agent:{self._config.agent_id}"
        cap_triples = "\n".join(
            f'  <{base_uri}> <mesh:capability> "{cap}" .'
            for cap in self._config.capabilities
        )
        sparql_update = (
            "PREFIX mesh: <https://meshpocalypse.ai/ontology#>\n\n"
            f"DELETE WHERE {{ <{base_uri}> <mesh:capability> ?c . }};\n\n"
            f"INSERT DATA {{\n{cap_triples}\n"
            f"  <{base_uri}> <mesh:trustLevel> {self._config.trust_level.value} .\n"
            "}}"
        )

        endpoint = f"{graph_url}/update"
        try:
            req = urllib.request.Request(
                endpoint,
                data=sparql_update.encode("utf-8"),
                headers={"Content-Type": "application/sparql-update"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=5) as resp:
                if resp.status >= 400:
                    logger.warning(
                        "[SwarmNode] Graph registration failed (HTTP %s) — continuing without graph",
                        resp.status,
                    )
        except Exception as exc:
            logger.warning(
                "[SwarmNode] Graph registration error — continuing without graph: %s", exc
            )
