"""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 socket
import time
import uuid
from collections.abc import Awaitable, Callable
from datetime import datetime
from typing import Any

from ana.audit import PublishAuditor
from ana.envelope import (
    Ack,
    Alert,
    Clock,
    Discovery,
    Envelope,
    EnvelopeError,
    Pulse,
    Query,
    Reply,
    parse_envelope,
)
from ana.subjects import BROADCAST, 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", "envelope")

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

    async def reply(
        self,
        data: Any,
        *,
        note: str | None = None,
        topic: str | None = None,
        on_subject: str | None = None,
    ) -> str:
        """Reply to the inbound message.

        Subject precedence (v0.5+):

        1. ``on_subject=`` kwarg, if explicitly passed — wins outright.
        2. The inbound ``Query.reply_to`` field, if the inbound envelope was
           a Query that named one. This honors the caller's wish that the
           reply land on a subject *they* are subscribed to.
        3. Default: ``<prefix>.<this-agent>.reply.<topic>`` — the v0.4
           behavior, kept for back-compat when the caller did not specify.

        The ``reply_for`` envelope field always carries the original inbound
        subject for correlation, independent of the publish subject.
        """
        parsed = self.bus.scheme.parse(self.subject)
        topic_ = topic if topic is not None else (parsed.topic or "default")

        inbound_reply_to: str | None = None
        if isinstance(self.envelope, Query) and self.envelope.reply_to:
            inbound_reply_to = self.envelope.reply_to

        publish_subject = on_subject or inbound_reply_to or self.bus.scheme.reply(
            self.bus.identity, topic_
        )

        env = Reply(
            **{"from": self.bus.identity},
            reply_for=self.subject,
            in_reply_to=self.envelope.request_id if isinstance(self.envelope, Query) else None,
            request_id=self.envelope.request_id if isinstance(self.envelope, Query) else None,
            data=data,
            note=note,
        )
        await self.bus.publish_envelope(publish_subject, env)
        return publish_subject

    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,
        host: str | None = None,
        stamp_clock: bool = True,
    ) -> None:
        if identity == BROADCAST:
            raise ValueError(
                f"identity={BROADCAST!r} is reserved for the broadcast pseudo-identity; pick a real name"
            )
        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] = []
        self.host = host if host is not None else socket.gethostname()
        self.stamp_clock = stamp_clock
        self._started_at = time.monotonic()
        self._pulse_task: asyncio.Task | None = None

    # ------------------------------------------------------------------
    # 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:
        if self._pulse_task is not None:
            self._pulse_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._pulse_task
            self._pulse_task = 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

    def _now_clock(self) -> Clock:
        return Clock(
            local_time=datetime.now().astimezone().isoformat(timespec="seconds"),
            uptime_s=int(time.monotonic() - self._started_at),
        )

    @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")
        # Auto-stamp clock + host on every outbound envelope unless caller
        # already set them. This is the v0.2 time/state-propagation hook —
        # receivers can use the trio (clock, host, ts) to anchor freshness
        # without relying on conversation history.
        payload = envelope.model_dump(by_alias=True, exclude_none=True)
        if self.stamp_clock and "clock" not in payload:
            payload["clock"] = self._now_clock().model_dump(exclude_none=True)
        if "host" not in payload:
            payload["host"] = self.host
        import json as _json

        body = _json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        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,
        on_subject: str | None = None,
    ) -> str:
        """Publish a Reply envelope.

        Subject precedence:

        - ``on_subject``, if passed, wins.
        - Otherwise default to ``<prefix>.<this-agent>.reply.<topic>``.

        Use ``on_subject=<query.reply_to>`` to honor a caller's explicit
        reply-routing hint. ``InboundContext.reply()`` does this
        automatically when the inbound envelope was a Query carrying
        ``reply_to``.
        """
        env = Reply(
            **{"from": self.identity},
            reply_for=reply_for,
            data=data,
            note=note,
            in_reply_to=in_reply_to,
            request_id=request_id,
        )
        subj = on_subject or 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

    async def pulse(
        self,
        *,
        activity: str | None = None,
        state: str | None = None,
        session_id: str | None = None,
        recent: list[str] | None = None,
        capabilities: dict[str, Any] | None = None,
    ) -> str:
        """Broadcast a Pulse envelope to ``<prefix>.<identity>.pulse``.

        Periodic emission of these is how peers learn your current
        time/state without you replying to a query — useful when an agent
        is doing autonomous work and wants to advertise progress.
        """
        env = Pulse(
            **{"from": self.identity},
            activity=activity,
            state=state,
            session_id=session_id,
            recent=recent or [],
            capabilities=capabilities or {},
        )
        subj = self.scheme.pulse(self.identity)
        await self.publish_envelope(subj, env)
        return subj

    async def alert(
        self,
        event: str,
        *,
        level: str = "info",
        data: Any = None,
        origin_subject: str | None = None,
    ) -> str:
        """Publish an Alert envelope to ``<prefix>.<identity>.alert.<event>``.

        One-way fire-and-forget notification — no reply expected. Receivers
        must be subscribed to this identity's namespace to see it; use
        :py:meth:`broadcast_alert` instead when you want every peer to see
        the alert without addressing each one.
        """
        env = Alert(
            **{"from": self.identity},
            event=event,
            level=level,  # type: ignore[arg-type]
            data=data,
            origin_subject=origin_subject,
        )
        subj = self.scheme.alert(self.identity, event)
        await self.publish_envelope(subj, env)
        return subj

    async def broadcast_alert(
        self,
        event: str,
        *,
        level: str = "info",
        data: Any = None,
        origin_subject: str | None = None,
    ) -> str:
        """Publish an Alert on the reserved ``broadcast`` pseudo-identity.

        Subject: ``<prefix>.broadcast.alert.<event>``. ana-listener daemons
        with the default config also subscribe to ``<prefix>.broadcast.>``,
        so this single publish reaches the whole fleet — fixes the v0.3
        gap where alerts on a sender's own namespace were invisible to
        scope=self subscribers.

        The envelope's ``from`` field still carries the real sender
        identity; only the subject uses the reserved pseudo-identity.
        """
        env = Alert(
            **{"from": self.identity},
            event=event,
            level=level,  # type: ignore[arg-type]
            data=data,
            origin_subject=origin_subject,
        )
        subj = self.scheme.broadcast("alert", event)
        await self.publish_envelope(subj, env)
        return subj

    async def subscribe_broadcasts(self, handler: EnvelopeHandler) -> Any:
        """Subscribe to ``<prefix>.broadcast.>``.

        Listeners that want to see fleet-wide broadcasts in addition to
        their own ``<prefix>.<identity>.>`` scope call this on top of
        :py:meth:`subscribe_to_me`. The CLI daemon does it by default
        unless ``--no-subscribe-broadcast`` is passed.
        """
        return await self.subscribe(self.scheme.all_broadcasts(), handler)

    async def start_pulse_loop(
        self,
        interval_s: float,
        *,
        activity_fn: Callable[[], str | None] | None = None,
        state_fn: Callable[[], str | None] | None = None,
        capabilities_fn: Callable[[], dict[str, Any]] | None = None,
    ) -> None:
        """Spawn a background task that emits a Pulse every ``interval_s``.

        ``activity_fn`` / ``state_fn`` / ``capabilities_fn`` are called
        synchronously at each tick so the daemon can surface fresh info.
        Pass ``None`` for any of them to leave the corresponding field empty.

        Cancelled automatically by :py:meth:`close`. Idempotent — calling
        twice replaces the previous loop.
        """
        if self._pulse_task is not None:
            self._pulse_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._pulse_task

        async def _loop() -> None:
            while True:
                try:
                    await self.pulse(
                        activity=activity_fn() if activity_fn else None,
                        state=state_fn() if state_fn else None,
                        capabilities=capabilities_fn() if capabilities_fn else None,
                    )
                except Exception:
                    log.exception("pulse loop publish failed")
                await asyncio.sleep(interval_s)

        self._pulse_task = asyncio.create_task(_loop(), name=f"ana-pulse-{self.identity}")

    # ------------------------------------------------------------------
    # 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:
            try:
                env = parse_envelope(msg.data)
            except EnvelopeError as e:
                log.warning("envelope parse failed on %s: %s", msg.subject, e)
                return
            ctx = InboundContext(subject=msg.subject, bus=self, raw=msg.data, envelope=env)
            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)
        responder_reply_subject = self.scheme.reply(to, topic)
        caller_reply_subject = self.scheme.reply(self.identity, 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)

        # v0.5: subscribe to BOTH the responder's reply subject (pre-v0.5 path)
        # AND our own reply subject (v0.5+ responders honor query.reply_to and
        # publish there).
        subs: list[Any] = [
            await self.subscribe(responder_reply_subject, _on_reply),
            await self.subscribe(caller_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)
