"""AgentBus — the async client.

Thin layer on top of ``nats-py`` that:

- Builds and parses envelopes via ``ana.envelope``
- Names subjects via ``SubjectScheme`` (configurable per fleet)
- Routes inbound messages to handlers registered with ``on()``, ``on_query()``,
  ``on_reply()``, ``on_ack()`` or via direct ``subscribe()``
- Audits every publish via ``PublishAuditor``
- Provides ``query_and_wait()`` for the request/reply pattern (subscribes to
  the reply subject, publishes the query, awaits the first matching reply)

Connection details are intentionally permissive: anonymous core-NATS is the
default since that's what most lightweight A2A buses run as. Token / user-pass
auth surfaces through ``nats_options``.

Core-NATS gotcha: a publish made while the target has no active subscription
is **lost forever** (no JetStream replay). When that matters, use
``query_and_wait()`` (it sets up the reply sub before publishing) or coordinate
sub-startup out of band.
"""
from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import uuid
from collections.abc import Awaitable, Callable
from typing import Any

from ana.audit import PublishAuditor
from ana.envelope import (
    Ack,
    Discovery,
    Envelope,
    EnvelopeError,
    Query,
    Reply,
    parse_envelope,
)
from ana.subjects import SubjectScheme

log = logging.getLogger("ana")

EnvelopeHandler = Callable[[Envelope, "InboundContext"], Awaitable[None]]


class InboundContext:
    """Passed to handlers alongside the parsed envelope.

    Holds the raw subject the message arrived on plus a back-reference to the
    bus so handlers can reply / ack without juggling state themselves.
    """

    __slots__ = ("subject", "bus", "raw")

    def __init__(self, subject: str, bus: AgentBus, raw: bytes) -> None:
        self.subject = subject
        self.bus = bus
        self.raw = raw

    async def reply(self, data: Any, *, note: str | None = None, topic: str | None = None) -> str:
        """Reply on the conventional reply subject derived from the inbound one.

        If the inbound subject is ``<prefix>.<inbound-agent>.query.<topic>``,
        the reply lands on ``<prefix>.<this-agent>.reply.<topic>``. The
        ``reply_for`` field carries the original inbound subject for the
        receiver's correlation logic.
        """
        parsed = self.bus.scheme.parse(self.subject)
        topic_ = topic if topic is not None else (parsed.topic or "default")
        env = Reply(
            **{"from": self.bus.identity},
            reply_for=self.subject,
            data=data,
            note=note,
        )
        subj = self.bus.scheme.reply(self.bus.identity, topic_)
        await self.bus.publish_envelope(subj, env)
        return subj

    async def ack(self, *, note: str | None = None) -> str:
        parsed = self.bus.scheme.parse(self.subject)
        topic_ = parsed.topic or "default"
        env = Ack(**{"from": self.bus.identity}, ack_for=self.subject, note=note)
        subj = self.bus.scheme.ack(self.bus.identity, topic_)
        await self.bus.publish_envelope(subj, env)
        return subj


class AgentBus:
    def __init__(
        self,
        identity: str,
        *,
        nats_url: str = "nats://127.0.0.1:4222",
        scheme: SubjectScheme | None = None,
        connection_name: str | None = None,
        audit_log_path: str | os.PathLike | None = None,
        nats_options: dict[str, Any] | None = None,
    ) -> None:
        self.identity = identity
        self.nats_url = nats_url
        self.scheme = scheme or SubjectScheme()
        self.connection_name = connection_name or f"{identity}-agent"
        self.auditor = PublishAuditor(log_path=audit_log_path)
        self._nats_options = nats_options or {}
        self._nc: Any = None
        self._subs: list[Any] = []

    # ------------------------------------------------------------------
    # lifecycle
    # ------------------------------------------------------------------
    async def __aenter__(self) -> AgentBus:
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.close()

    async def connect(self) -> None:
        import nats  # local import keeps `pip install ana[no-nats]` testable

        self._nc = await nats.connect(
            servers=[self.nats_url],
            name=self.connection_name,
            **self._nats_options,
        )

    async def close(self) -> None:
        for sub in self._subs:
            with contextlib.suppress(Exception):
                await sub.unsubscribe()
        self._subs.clear()
        if self._nc is not None:
            try:
                await self._nc.drain()
            finally:
                self._nc = None

    @property
    def is_connected(self) -> bool:
        return self._nc is not None and self._nc.is_connected

    # ------------------------------------------------------------------
    # publish
    # ------------------------------------------------------------------
    async def publish_envelope(self, subject: str, envelope: Envelope) -> None:
        if self._nc is None:
            raise RuntimeError("AgentBus not connected; call connect() first")
        body = envelope.to_bytes()
        self.auditor.track(subject, len(body), envelope.type)
        await self._nc.publish(subject, body)

    async def query(
        self,
        to: str,
        query: str,
        *,
        topic: str = "default",
        fields: list[str] | None = None,
        reply_to: str | None = None,
        request_id: str | None = None,
    ) -> str:
        env = Query(
            **{"from": self.identity},
            to=to,
            query=query,
            fields=fields,
            reply_to=reply_to or self.scheme.reply(self.identity, topic),
            request_id=request_id or uuid.uuid4().hex[:12],
        )
        subj = self.scheme.query(to, topic)
        await self.publish_envelope(subj, env)
        return subj

    async def reply(
        self,
        *,
        reply_for: str,
        data: Any,
        topic: str = "default",
        note: str | None = None,
        in_reply_to: str | None = None,
        request_id: str | None = None,
    ) -> str:
        env = Reply(
            **{"from": self.identity},
            reply_for=reply_for,
            data=data,
            note=note,
            in_reply_to=in_reply_to,
            request_id=request_id,
        )
        subj = self.scheme.reply(self.identity, topic)
        await self.publish_envelope(subj, env)
        return subj

    async def ack(
        self,
        *,
        ack_for: str,
        topic: str = "default",
        note: str | None = None,
        received_subjects: list[str] | None = None,
    ) -> str:
        env = Ack(
            **{"from": self.identity},
            ack_for=ack_for,
            note=note,
            received_subjects=received_subjects,
        )
        subj = self.scheme.ack(self.identity, topic)
        await self.publish_envelope(subj, env)
        return subj

    async def announce_discovery(
        self,
        role: str,
        *,
        cluster: str | None = None,
        subjects_owned: list[str] | None = None,
        capabilities: dict[str, Any] | None = None,
    ) -> str:
        env = Discovery(
            **{"from": self.identity},
            role=role,
            cluster=cluster,
            subjects_owned=subjects_owned or [self.scheme.all_for_agent(self.identity)],
            capabilities=capabilities or {},
        )
        subj = self.scheme.discovery(self.identity)
        await self.publish_envelope(subj, env)
        return subj

    # ------------------------------------------------------------------
    # subscribe
    # ------------------------------------------------------------------
    async def subscribe(
        self,
        subject: str,
        handler: EnvelopeHandler,
        *,
        queue: str | None = None,
    ) -> Any:
        if self._nc is None:
            raise RuntimeError("AgentBus not connected; call connect() first")

        async def _cb(msg: Any) -> None:
            ctx = InboundContext(subject=msg.subject, bus=self, raw=msg.data)
            try:
                env = parse_envelope(msg.data)
            except EnvelopeError as e:
                log.warning("envelope parse failed on %s: %s", msg.subject, e)
                return
            try:
                await handler(env, ctx)
            except Exception:
                log.exception("handler crashed on %s", msg.subject)

        sub = await self._nc.subscribe(subject, cb=_cb, queue=queue)
        self._subs.append(sub)
        return sub

    async def subscribe_all(self, handler: EnvelopeHandler) -> Any:
        """Subscribe to everything under the scheme's prefix (``<prefix>.>``)."""
        return await self.subscribe(self.scheme.all_under_prefix(), handler)

    async def subscribe_to_me(self, handler: EnvelopeHandler) -> Any:
        """Subscribe to everything addressed under this agent's subject space."""
        return await self.subscribe(self.scheme.all_for_agent(self.identity), handler)

    # ------------------------------------------------------------------
    # request/reply convenience
    # ------------------------------------------------------------------
    async def query_and_wait(
        self,
        to: str,
        query: str,
        *,
        topic: str = "default",
        fields: list[str] | None = None,
        timeout_s: float = 10.0,
        accept_ack: bool = False,
    ) -> Reply | Ack | None:
        """Publish a query, await the first matching reply (or ack), return it.

        Subscribes to the responder's reply (and optionally ack) subject BEFORE
        publishing, so the core-NATS "lost-if-no-sub" gotcha doesn't bite here.
        Returns ``None`` on timeout — callers decide whether that means retry,
        escalate, or give up.

        With ``accept_ack=True``, a matching ``Ack`` envelope from the target
        also satisfies the wait. Useful for liveness probes against
        listener-only daemons that don't publish replies.
        """
        request_id = uuid.uuid4().hex[:12]
        query_subject = self.scheme.query(to, topic)
        reply_subject = self.scheme.reply(to, topic)
        ack_subject = self.scheme.ack(to, topic)
        future: asyncio.Future[Reply | Ack] = asyncio.get_event_loop().create_future()

        async def _on_reply(env: Envelope, _ctx: InboundContext) -> None:
            if future.done():
                return
            if isinstance(env, Reply) and (
                env.request_id == request_id
                or env.in_reply_to == request_id
                or env.reply_for == query_subject
            ):
                future.set_result(env)

        async def _on_ack(env: Envelope, _ctx: InboundContext) -> None:
            if future.done():
                return
            if isinstance(env, Ack) and env.ack_for == query_subject:
                future.set_result(env)

        subs: list[Any] = [await self.subscribe(reply_subject, _on_reply)]
        if accept_ack:
            subs.append(await self.subscribe(ack_subject, _on_ack))
        try:
            await self.query(
                to,
                query,
                topic=topic,
                fields=fields,
                request_id=request_id,
            )
            try:
                return await asyncio.wait_for(future, timeout=timeout_s)
            except asyncio.TimeoutError:
                return None
        finally:
            for sub in subs:
                with contextlib.suppress(Exception):
                    await sub.unsubscribe()
                if sub in self._subs:
                    self._subs.remove(sub)
