"""Event-driven swarm runtime built on top of the persistent event bus."""
from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass
from typing import Any, Awaitable, Callable
from uuid import uuid4

from logos_adk.checkpoint_json import JSONObject, ensure_json_object
from logos_adk.event_bus import EventBus, EventRecord
from logos_adk.runnable import Runnable
from logos_adk.types import Response, Usage
from logos_adk.ui_events import make_ui_event


SwarmTarget = Runnable | Callable[[EventRecord], Awaitable[Any]]
PublishTopicResolver = str | Callable[[EventRecord, JSONObject], str]


@dataclass
class SwarmSubscription:
    """A topic subscription handled by a runnable or async function."""

    topic: str
    handler_name: str
    target: SwarmTarget
    publish_to: PublishTopicResolver | None = None


class Swarm:
    """Claim events from a bus, process them, and optionally publish derived events."""

    def __init__(
        self,
        name: str,
        *,
        event_bus: EventBus,
        retry_delay_seconds: float = 0.0,
        lease_seconds: float = 60.0,
        slow_step_notice_seconds: float = 12.0,
        slow_step_repeat_seconds: float = 15.0,
    ) -> None:
        self._name = name
        self._event_bus = event_bus
        self._retry_delay_seconds = retry_delay_seconds
        self._lease_seconds = lease_seconds
        self._slow_step_notice_seconds = max(0.01, float(slow_step_notice_seconds))
        self._slow_step_repeat_seconds = max(0.01, float(slow_step_repeat_seconds))
        self._subscriptions: list[SwarmSubscription] = []

    @property
    def name(self) -> str:
        return self._name

    @property
    def event_bus(self) -> EventBus:
        return self._event_bus

    def subscribe(
        self,
        topic: str,
        target: SwarmTarget,
        *,
        name: str | None = None,
        publish_to: PublishTopicResolver | None = None,
    ) -> Swarm:
        handler_name = name or getattr(target, "name", None) or getattr(target, "__name__", None) or topic
        self._subscriptions.append(
            SwarmSubscription(
                topic=topic,
                handler_name=str(handler_name),
                target=target,
                publish_to=publish_to,
            )
        )
        return self

    async def drain(self, *, max_events: int | None = None) -> int:
        """Process claimable events until the backlog is empty or max_events is reached.

        Handler errors are caught and logged — a failing handler for one
        subscription does not block processing of other subscriptions.
        The failed event is requeued by ``_process_event`` with a delay
        that prevents immediate re-claim within the same drain cycle.
        """
        processed = 0
        _failed: dict[str, str] = {}  # event_id → error message
        while True:
            made_progress = False
            for subscription in self._subscriptions:
                if max_events is not None and processed >= max_events:
                    return processed
                event = self._event_bus.claim(
                    self._consumer_name(subscription),
                    [subscription.topic],
                    lease_seconds=self._lease_seconds,
                )
                if event is None:
                    continue
                if event.id in _failed:
                    # Already failed in this drain cycle — skip to avoid infinite retry.
                    # Release the claim so it can be picked up on next drain call.
                    self._event_bus.requeue(
                        event.id,
                        self._consumer_name(subscription),
                        delay_seconds=max(self._retry_delay_seconds, 1.0),
                        error=_failed[event.id],
                    )
                    continue
                made_progress = True
                try:
                    await self._process_event(subscription, event)
                except Exception as exc:
                    import logging
                    logging.getLogger("logos_adk.swarm").error(
                        "Handler %s failed for event %s [%s]: %s",
                        subscription.handler_name,
                        event.id,
                        event.topic,
                        exc,
                    )
                    _failed[event.id] = f"{type(exc).__name__}: {exc}"
                processed += 1
            if not made_progress:
                return processed

    async def drain_once(self) -> int:
        """Process at most one event."""
        return await self.drain(max_events=1)

    async def drain_reactive(
        self,
        *,
        timeout: float = 5.0,
        max_events: int = 20,
    ) -> int:
        """Wait for a NOTIFY wake-up, then drain available events.

        Uses PostgreSQL ``LISTEN/NOTIFY`` to avoid busy-polling.
        Falls back to a plain ``drain()`` after *timeout* seconds if no
        notification arrives (ensures liveness even without NOTIFY).

        This is a drop-in replacement for the pattern::

            await swarm.drain(max_events=N)
            await asyncio.sleep(interval)

        Instead, callers simply ``await swarm.drain_reactive()``.
        The method returns immediately after draining all available events.

        Requires the event bus to be a :class:`PostgreSQLEventBus`.
        For other backends, silently falls back to sleep + drain.
        """
        from logos_adk.event_bus import PostgreSQLEventBus

        bus = self._event_bus
        if isinstance(bus, PostgreSQLEventBus):
            # Wait for a push notification (or timeout)
            await bus.wait_for_notify(timeout=timeout)
        else:
            # Non-PostgreSQL fallback: simple sleep like the old behaviour
            await asyncio.sleep(timeout)
        return await self.drain(max_events=max_events)

    async def dispatch(
        self,
        prompt: str,
        *,
        request_topic: str,
        response_topic: str | None = None,
        payload: JSONObject | None = None,
        session_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        campaign_id: str | None = None,
        publisher: str | None = None,
        max_events: int | None = None,
    ) -> Response:
        """Publish a request into the swarm and drain matching work immediately."""
        effective_correlation_id = correlation_id or session_id or uuid4().hex
        request_payload = dict(payload or {})
        if not isinstance(request_payload.get("prompt"), str) or not request_payload["prompt"].strip():
            request_payload["prompt"] = prompt
        request = self._event_bus.publish(
            request_topic,
            ensure_json_object(request_payload, path=f"swarm.dispatch[{request_topic}]"),
            campaign_id=campaign_id,
            correlation_id=effective_correlation_id,
            run_id=run_id,
            publisher=publisher or self._name,
        )
        await self.drain(max_events=max_events)
        request_after = self._event_bus.get_event(request.id)
        if response_topic is not None:
            responses = self._event_bus.list_events(
                topic=response_topic,
                correlation_id=effective_correlation_id,
                run_id=run_id,
            )
            if responses:
                latest = responses[-1]
                return Response(
                    content=_payload_to_text(latest.payload),
                    tool_calls=[],
                    usage=Usage(0, 0),
                    raw={
                        "swarm_dispatch": {
                            "swarm": self._name,
                            "request_topic": request_topic,
                            "response_topic": response_topic,
                            "request_event_id": request.id,
                            "response_event_id": latest.id,
                        },
                        "event_payload": dict(latest.payload),
                    },
                )
        request_status = request_after.status if request_after is not None else "queued"
        return Response(
            content=f"Swarm request status: {request_status}",
            tool_calls=[],
            usage=Usage(0, 0),
            raw={
                "swarm_dispatch": {
                    "swarm": self._name,
                    "request_topic": request_topic,
                    "response_topic": response_topic,
                    "request_event_id": request.id,
                    "request_status": request_status,
                },
                "event_payload": dict(request.payload),
            },
            status="completed" if request_status == "completed" else "paused",
        )

    async def dispatch_events(
        self,
        prompt: str,
        *,
        request_topic: str,
        response_topic: str | None = None,
        payload: JSONObject | None = None,
        session_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        campaign_id: str | None = None,
        publisher: str | None = None,
        max_events: int | None = None,
        trace_id: str | None = None,
    ):
        """Dispatch into the swarm and yield real workflow events as work is processed."""
        effective_run_id = run_id or uuid4().hex
        effective_correlation_id = correlation_id or session_id or effective_run_id
        sequence = 0
        yield make_ui_event(
            event="graph_started",
            event_type="swarm",
            graph_name=self._name,
            graph_run_id=effective_run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self._name,
            status="running",
            output={"request_topic": request_topic, "response_topic": response_topic},
        )
        sequence += 1

        request_payload = dict(payload or {})
        if not isinstance(request_payload.get("prompt"), str) or not request_payload["prompt"].strip():
            request_payload["prompt"] = prompt
        request = self._event_bus.publish(
            request_topic,
            ensure_json_object(request_payload, path=f"swarm.dispatch[{request_topic}]"),
            campaign_id=campaign_id,
            correlation_id=effective_correlation_id,
            run_id=effective_run_id,
            publisher=publisher or self._name,
        )

        processed = 0
        while True:
            made_progress = False
            for subscription in self._subscriptions:
                if max_events is not None and processed >= max_events:
                    break
                event = self._event_bus.claim(
                    self._consumer_name(subscription),
                    [subscription.topic],
                    lease_seconds=self._lease_seconds,
                )
                if event is None:
                    continue
                made_progress = True
                yield make_ui_event(
                    event="node_started",
                    event_type="swarm",
                    graph_name=self._name,
                    graph_run_id=effective_run_id,
                    sequence=sequence,
                    session_id=session_id,
                    correlation_id=effective_correlation_id,
                    trace_id=trace_id,
                    node_name=subscription.handler_name,
                    status="running",
                    route_source=event.topic,
                    topic=event.topic,
                    publisher=event.publisher,
                    output=event.payload,
                )
                sequence += 1
                handler_task = asyncio.create_task(self._invoke_handler(subscription.target, event))
                started = time.monotonic()
                notices_sent = 0
                try:
                    while True:
                        wait_seconds = (
                            self._slow_step_notice_seconds
                            if notices_sent == 0
                            else self._slow_step_repeat_seconds
                        )
                        done, _ = await asyncio.wait({handler_task}, timeout=wait_seconds)
                        if handler_task in done:
                            result_payload = await handler_task
                            break

                        notices_sent += 1
                        waited_seconds = int(max(1, round(time.monotonic() - started)))
                        yield make_ui_event(
                            event="node_progress",
                            event_type="swarm",
                            graph_name=self._name,
                            graph_run_id=effective_run_id,
                            sequence=sequence,
                            session_id=session_id,
                            correlation_id=effective_correlation_id,
                            trace_id=trace_id,
                            node_name=subscription.handler_name,
                            status="running",
                            route_source=event.topic,
                            topic=event.topic,
                            publisher=event.publisher,
                            output={
                                "waited_seconds": waited_seconds,
                                "message": (
                                    "Шаг выполняется дольше обычного. "
                                    f"Ждём ответ сервиса уже {waited_seconds} сек."
                                ),
                            },
                            metadata={
                                "user_facing": {
                                    "kind": "step",
                                    "title": _humanize_label(subscription.handler_name),
                                    "detail": (
                                        "Шаг выполняется дольше обычного. "
                                        f"Ждём ответ сервиса уже {waited_seconds} сек."
                                    ),
                                    "step_key": f"step:{subscription.handler_name}",
                                    "status": "running",
                                    "visible": True,
                                }
                            },
                        )
                        sequence += 1
                except Exception as exc:
                    if not handler_task.done():
                        handler_task.cancel()
                    self._event_bus.requeue(
                        event.id,
                        self._consumer_name(subscription),
                        delay_seconds=self._retry_delay_seconds,
                        error=f"{type(exc).__name__}: {exc}",
                    )
                    raise
                finally:
                    if not handler_task.done():
                        handler_task.cancel()
                publish_topic = None
                if subscription.publish_to is not None:
                    publish_topic = self._resolve_publish_topic(subscription.publish_to, event, result_payload)
                    self._event_bus.publish(
                        publish_topic,
                        result_payload,
                        correlation_id=event.correlation_id,
                        run_id=event.run_id,
                        publisher=subscription.handler_name,
                    )
                self._event_bus.ack(event.id, self._consumer_name(subscription))
                yield make_ui_event(
                    event="node_completed",
                    event_type="swarm",
                    graph_name=self._name,
                    graph_run_id=effective_run_id,
                    sequence=sequence,
                    session_id=session_id,
                    correlation_id=effective_correlation_id,
                    trace_id=trace_id,
                    node_name=subscription.handler_name,
                    status="completed",
                    route_source=event.topic,
                    route_target=publish_topic,
                    topic=publish_topic,
                    publisher=subscription.handler_name,
                    output=result_payload,
                )
                sequence += 1
                processed += 1
            if not made_progress or (max_events is not None and processed >= max_events):
                break

        request_after = self._event_bus.get_event(request.id)
        if response_topic is not None:
            responses = self._event_bus.list_events(
                topic=response_topic,
                correlation_id=effective_correlation_id,
                run_id=effective_run_id,
            )
            if responses:
                latest = responses[-1]
                response = Response(
                    content=_payload_to_text(latest.payload),
                    tool_calls=[],
                    usage=Usage(0, 0),
                    raw={
                        "swarm_dispatch": {
                            "swarm": self._name,
                            "request_topic": request_topic,
                            "response_topic": response_topic,
                            "request_event_id": request.id,
                            "response_event_id": latest.id,
                        },
                        "event_payload": dict(latest.payload),
                    },
                )
            else:
                request_status = request_after.status if request_after is not None else "queued"
                response = Response(
                    content=f"Swarm request status: {request_status}",
                    tool_calls=[],
                    usage=Usage(0, 0),
                    raw={
                        "swarm_dispatch": {
                            "swarm": self._name,
                            "request_topic": request_topic,
                            "response_topic": response_topic,
                            "request_event_id": request.id,
                            "request_status": request_status,
                        },
                        "event_payload": dict(request.payload),
                    },
                    status="completed" if request_status == "completed" else "paused",
                )
        else:
            request_status = request_after.status if request_after is not None else "queued"
            response = Response(
                content=f"Swarm request status: {request_status}",
                tool_calls=[],
                usage=Usage(0, 0),
                raw={
                    "swarm_dispatch": {
                        "swarm": self._name,
                        "request_topic": request_topic,
                        "response_topic": response_topic,
                        "request_event_id": request.id,
                        "request_status": request_status,
                    },
                    "event_payload": dict(request.payload),
                },
                status="completed" if request_status == "completed" else "paused",
            )
        yield make_ui_event(
            event="graph_completed",
            event_type="swarm",
            graph_name=self._name,
            graph_run_id=effective_run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self._name,
            status=response.status,
            output=response,
            route_source=request.topic,
            route_target=response_topic,
        )

    async def _process_event(self, subscription: SwarmSubscription, event: EventRecord) -> None:
        consumer = self._consumer_name(subscription)
        try:
            result_payload = await self._invoke_handler(subscription.target, event)
            if subscription.publish_to is not None:
                response_topic = self._resolve_publish_topic(subscription.publish_to, event, result_payload)
                self._event_bus.publish(
                    response_topic,
                    result_payload,
                    correlation_id=event.correlation_id,
                    run_id=event.run_id,
                    publisher=subscription.handler_name,
                )
            self._event_bus.ack(event.id, consumer)
        except Exception as exc:
            self._event_bus.requeue(
                event.id,
                consumer,
                delay_seconds=self._retry_delay_seconds,
                error=f"{type(exc).__name__}: {exc}",
            )
            raise

    def _consumer_name(self, subscription: SwarmSubscription) -> str:
        return f"{self._name}:{subscription.handler_name}"

    def _resolve_publish_topic(
        self,
        publish_to: PublishTopicResolver,
        event: EventRecord,
        payload: JSONObject,
    ) -> str:
        if callable(publish_to):
            topic = publish_to(event, payload)
        else:
            topic = publish_to
        if not isinstance(topic, str) or not topic.strip():
            raise ValueError("Swarm publish_to must resolve to a non-empty topic")
        return topic.strip()

    async def _invoke_handler(self, target: SwarmTarget, event: EventRecord) -> JSONObject:
        if isinstance(target, Runnable):
            prompt = _prompt_from_payload(event.payload)
            response = await target.run(
                prompt,
                session_id=event.correlation_id,
                correlation_id=event.correlation_id,
            )
            return _response_to_payload(response)
        result = await target(event)
        return _coerce_payload(result)

def _humanize_label(value: str) -> str:
    return value.replace("_", " ").replace("-", " ").strip().title()


def _prompt_from_payload(payload: JSONObject) -> str:
    prompt = payload.get("prompt")
    if isinstance(prompt, str) and prompt:
        return prompt
    return json.dumps(payload, ensure_ascii=False)


def _response_to_payload(response: Response) -> JSONObject:
    payload: JSONObject = {"content": response.content}
    if response.raw:
        payload["raw"] = ensure_json_object(dict(response.raw), path="swarm.response.raw")
    return payload


def _coerce_payload(value: Any) -> JSONObject:
    if isinstance(value, Response):
        return _response_to_payload(value)
    if value is None:
        return {"status": "ok"}
    if isinstance(value, dict):
        return ensure_json_object(value, path="swarm.handler.payload")
    return {"output": str(value)}


def _payload_to_text(payload: JSONObject) -> str:
    content = payload.get("content")
    if isinstance(content, str) and content.strip():
        return content
    output = payload.get("output")
    if isinstance(output, str) and output.strip():
        return output
    return json.dumps(payload, ensure_ascii=False)


__all__ = [
    "Swarm",
    "SwarmSubscription",
]
