"""Scalability primitives for Logos ADK."""
from __future__ import annotations

import asyncio
from collections import defaultdict
from dataclasses import asdict, dataclass
from datetime import datetime, UTC
from pathlib import Path
import json
import random
import sqlite3
from typing import TYPE_CHECKING, Any, Awaitable, Callable
from uuid import uuid4

from logos_adk.postgres_utils import import_psycopg, resolve_database_url
from logos_adk.types import Response, ToolCall, Usage

if TYPE_CHECKING:
    from logos_adk.agent import Agent


@dataclass
class TaskRecord:
    """A queued task and its execution state."""

    id: str
    agent_name: str
    status: str
    created_at: datetime
    started_at: datetime | None = None
    finished_at: datetime | None = None
    error: str | None = None
    idempotency_key: str | None = None
    attempt: int = 0
    replay_of: str | None = None
    response_json: str | None = None


def _serialize_response(response: Response) -> str:
    return json.dumps(asdict(response), sort_keys=True)


def _deserialize_response(payload: str) -> Response:
    data = json.loads(payload)
    usage = data.get("usage") or {}
    tool_calls = [
        ToolCall(
            id=item.get("id", ""),
            name=item.get("name", ""),
            arguments=item.get("arguments") or {},
        )
        for item in (data.get("tool_calls") or [])
    ]
    return Response(
        content=data.get("content", ""),
        tool_calls=tool_calls,
        usage=Usage(
            input_tokens=int(usage.get("input_tokens", 0)),
            output_tokens=int(usage.get("output_tokens", 0)),
        ),
        raw=data.get("raw") or {},
        reasoning=data.get("reasoning"),
    )


class SQLiteTaskQueue:
    """Persistent async task queue backed by SQLite."""

    def __init__(
        self,
        path: str = ".logos_adk/tasks.db",
        concurrency: int = 4,
        *,
        idempotency_ttl_seconds: int = 86_400,
    ) -> None:
        self.path = path
        self.concurrency = concurrency
        self.idempotency_ttl_seconds = idempotency_ttl_seconds
        self._semaphore = asyncio.Semaphore(concurrency)
        self._schema_ready = False

    def _connect(self) -> sqlite3.Connection:
        db_path = Path(self.path)
        if db_path.parent and not db_path.parent.exists():
            db_path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(self.path)
        connection.row_factory = sqlite3.Row
        return connection

    def _ensure_schema(self, connection: sqlite3.Connection) -> None:
        if self._schema_ready:
            return
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS task_queue (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                status TEXT NOT NULL,
                created_at TEXT NOT NULL,
                started_at TEXT,
                finished_at TEXT,
                error TEXT,
                idempotency_key TEXT,
                attempt INTEGER NOT NULL DEFAULT 0,
                replay_of TEXT,
                response_json TEXT
            )
            """
        )
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS task_queue_dead_letter (
                task_id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                failed_at TEXT NOT NULL,
                error TEXT,
                attempt INTEGER NOT NULL DEFAULT 0,
                replayable INTEGER NOT NULL DEFAULT 1
            )
            """
        )
        columns = {
            row["name"]
            for row in connection.execute("PRAGMA table_info(task_queue)").fetchall()
        }
        if "idempotency_key" not in columns:
            connection.execute("ALTER TABLE task_queue ADD COLUMN idempotency_key TEXT")
        if "attempt" not in columns:
            connection.execute("ALTER TABLE task_queue ADD COLUMN attempt INTEGER NOT NULL DEFAULT 0")
        if "replay_of" not in columns:
            connection.execute("ALTER TABLE task_queue ADD COLUMN replay_of TEXT")
        if "response_json" not in columns:
            connection.execute("ALTER TABLE task_queue ADD COLUMN response_json TEXT")
        connection.execute(
            "CREATE INDEX IF NOT EXISTS idx_task_queue_status_created ON task_queue (status, created_at)"
        )
        connection.execute(
            "CREATE INDEX IF NOT EXISTS idx_task_queue_idempotency ON task_queue (idempotency_key, created_at)"
        )
        connection.commit()
        self._schema_ready = True

    def _insert_task(self, record: TaskRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                """
                INSERT INTO task_queue (
                    id, agent_name, status, created_at, started_at, finished_at, error,
                    idempotency_key, attempt, replay_of, response_json
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    record.id,
                    record.agent_name,
                    record.status,
                    record.created_at.isoformat(),
                    record.started_at.isoformat() if record.started_at else None,
                    record.finished_at.isoformat() if record.finished_at else None,
                    record.error,
                    record.idempotency_key,
                    record.attempt,
                    record.replay_of,
                    record.response_json,
                ),
            )
            connection.commit()

    def _update_task(
        self,
        task_id: str,
        *,
        status: str,
        started_at: datetime | None = None,
        finished_at: datetime | None = None,
        error: str | None = None,
        attempt: int | None = None,
        response_json: str | None = None,
    ) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                """
                UPDATE task_queue
                SET status = ?, started_at = ?, finished_at = ?, error = ?,
                    attempt = COALESCE(?, attempt),
                    response_json = COALESCE(?, response_json)
                WHERE id = ?
                """,
                (
                    status,
                    started_at.isoformat() if started_at else None,
                    finished_at.isoformat() if finished_at else None,
                    error,
                    attempt,
                    response_json,
                    task_id,
                ),
            )
            connection.commit()

    def _lookup_idempotent(self, idempotency_key: str) -> Response | None:
        cutoff = datetime.now(UTC).timestamp() - self.idempotency_ttl_seconds
        with self._connect() as connection:
            self._ensure_schema(connection)
            rows = connection.execute(
                """
                SELECT response_json, created_at
                FROM task_queue
                WHERE idempotency_key = ? AND status = 'completed' AND response_json IS NOT NULL
                ORDER BY created_at DESC
                LIMIT 5
                """,
                (idempotency_key,),
            ).fetchall()
        for row in rows:
            created_at = datetime.fromisoformat(row["created_at"]).timestamp()
            if created_at >= cutoff and row["response_json"]:
                return _deserialize_response(row["response_json"])
        return None

    def _record_dead_letter(self, record: TaskRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                """
                INSERT OR REPLACE INTO task_queue_dead_letter
                (task_id, agent_name, failed_at, error, attempt, replayable)
                VALUES (?, ?, ?, ?, ?, 1)
                """,
                (
                    record.id,
                    record.agent_name,
                    datetime.now(UTC).isoformat(),
                    record.error,
                    record.attempt,
                ),
            )
            connection.commit()

    async def submit(
        self,
        agent_name: str,
        work: Callable[[], Awaitable[Response]],
        *,
        idempotency_key: str | None = None,
        max_retries: int = 0,
        replay_of: str | None = None,
    ) -> Response:
        """Persist and execute a task with bounded concurrency."""
        if idempotency_key:
            existing = self._lookup_idempotent(idempotency_key)
            if existing is not None:
                return existing
        task_id = str(uuid4())
        record = TaskRecord(
            id=task_id,
            agent_name=agent_name,
            status="queued",
            created_at=datetime.now(UTC),
            idempotency_key=idempotency_key,
            replay_of=replay_of,
        )
        self._insert_task(record)

        async with self._semaphore:
            attempt = 0
            while True:
                attempt += 1
                started_at = datetime.now(UTC)
                self._update_task(task_id, status="running", started_at=started_at, attempt=attempt)
                try:
                    response = await work()
                except Exception as exc:
                    if attempt <= max_retries:
                        self._update_task(
                            task_id,
                            status="queued",
                            started_at=started_at,
                            finished_at=None,
                            error=str(exc),
                            attempt=attempt,
                        )
                        continue
                    failed_record = TaskRecord(
                        id=task_id,
                        agent_name=agent_name,
                        status="failed",
                        created_at=record.created_at,
                        started_at=started_at,
                        finished_at=datetime.now(UTC),
                        error=str(exc),
                        idempotency_key=idempotency_key,
                        attempt=attempt,
                        replay_of=replay_of,
                    )
                    self._update_task(
                        task_id,
                        status="failed",
                        started_at=started_at,
                        finished_at=failed_record.finished_at,
                        error=str(exc),
                        attempt=attempt,
                    )
                    self._record_dead_letter(failed_record)
                    raise
                response_json = _serialize_response(response)
                self._update_task(
                    task_id,
                    status="completed",
                    started_at=started_at,
                    finished_at=datetime.now(UTC),
                    attempt=attempt,
                    response_json=response_json,
                )
                return response

    async def replay(
        self,
        task_id: str,
        agent_name: str,
        work: Callable[[], Awaitable[Response]],
        *,
        max_retries: int = 0,
    ) -> Response:
        return await self.submit(
            agent_name,
            work,
            max_retries=max_retries,
            replay_of=task_id,
        )

    def stats(self) -> dict[str, int]:
        """Return task counts grouped by status."""
        with self._connect() as connection:
            self._ensure_schema(connection)
            rows = connection.execute(
                "SELECT status, COUNT(*) AS count FROM task_queue GROUP BY status"
            ).fetchall()
            dead_letter = connection.execute(
                "SELECT COUNT(*) AS count FROM task_queue_dead_letter"
            ).fetchone()
        stats = {row["status"]: row["count"] for row in rows}
        stats["dead_letter"] = dead_letter["count"] if dead_letter else 0
        return stats


class PostgreSQLTaskQueue:
    """Persistent async task queue backed by PostgreSQL."""

    def __init__(
        self,
        dsn: str,
        *,
        table_name: str = "task_queue",
        concurrency: int = 4,
        idempotency_ttl_seconds: int = 86_400,
    ) -> None:
        self.dsn = dsn
        self.table_name = table_name
        self.concurrency = concurrency
        self.idempotency_ttl_seconds = idempotency_ttl_seconds
        self._semaphore = asyncio.Semaphore(concurrency)

    def _connect(self) -> Any:
        psycopg = import_psycopg()
        return psycopg.connect(self.dsn)

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    agent_name TEXT NOT NULL,
                    status TEXT NOT NULL,
                    created_at TIMESTAMPTZ NOT NULL,
                    started_at TIMESTAMPTZ,
                    finished_at TIMESTAMPTZ,
                    error TEXT,
                    idempotency_key TEXT,
                    attempt INTEGER NOT NULL DEFAULT 0,
                    replay_of TEXT,
                    response_json JSONB
                )
                """
            )
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name}_dead_letter (
                    task_id TEXT PRIMARY KEY,
                    agent_name TEXT NOT NULL,
                    failed_at TIMESTAMPTZ NOT NULL,
                    error TEXT,
                    attempt INTEGER NOT NULL DEFAULT 0,
                    replayable BOOLEAN NOT NULL DEFAULT TRUE
                )
                """
            )
            cursor.execute(
                f"ALTER TABLE {self.table_name} ADD COLUMN IF NOT EXISTS idempotency_key TEXT"
            )
            cursor.execute(
                f"ALTER TABLE {self.table_name} ADD COLUMN IF NOT EXISTS attempt INTEGER NOT NULL DEFAULT 0"
            )
            cursor.execute(
                f"ALTER TABLE {self.table_name} ADD COLUMN IF NOT EXISTS replay_of TEXT"
            )
            cursor.execute(
                f"ALTER TABLE {self.table_name} ADD COLUMN IF NOT EXISTS response_json JSONB"
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_status "
                f"ON {self.table_name} (status)"
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_idempotency_created "
                f"ON {self.table_name} (idempotency_key, created_at DESC)"
            )
        connection.commit()

    def _insert_task(self, record: TaskRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name}
                    (id, agent_name, status, created_at, started_at, finished_at, error,
                     idempotency_key, attempt, replay_of, response_json)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        record.id,
                        record.agent_name,
                        record.status,
                        record.created_at,
                        record.started_at,
                        record.finished_at,
                        record.error,
                        record.idempotency_key,
                        record.attempt,
                        record.replay_of,
                        json.loads(record.response_json) if record.response_json else None,
                    ),
                )
            connection.commit()

    def _update_task(
        self,
        task_id: str,
        *,
        status: str,
        started_at: datetime | None = None,
        finished_at: datetime | None = None,
        error: str | None = None,
        attempt: int | None = None,
        response_json: str | None = None,
    ) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    UPDATE {self.table_name}
                    SET status = %s, started_at = %s, finished_at = %s, error = %s,
                        attempt = COALESCE(%s, attempt),
                        response_json = COALESCE(%s, response_json)
                    WHERE id = %s
                    """,
                    (
                        status,
                        started_at,
                        finished_at,
                        error,
                        attempt,
                        json.loads(response_json) if response_json else None,
                        task_id,
                    ),
                )
            connection.commit()

    def _lookup_idempotent(self, idempotency_key: str) -> Response | None:
        cutoff = datetime.now(UTC).timestamp() - self.idempotency_ttl_seconds
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    SELECT response_json, created_at
                    FROM {self.table_name}
                    WHERE idempotency_key = %s AND status = 'completed' AND response_json IS NOT NULL
                    ORDER BY created_at DESC
                    LIMIT 5
                    """,
                    (idempotency_key,),
                )
                rows = cursor.fetchall()
        for response_json, created_at in rows:
            created_timestamp = created_at.timestamp() if hasattr(created_at, "timestamp") else 0
            if created_timestamp >= cutoff and response_json is not None:
                return _deserialize_response(json.dumps(response_json))
        return None

    def _record_dead_letter(self, record: TaskRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name}_dead_letter
                    (task_id, agent_name, failed_at, error, attempt, replayable)
                    VALUES (%s, %s, %s, %s, %s, TRUE)
                    ON CONFLICT (task_id) DO UPDATE SET
                        failed_at = EXCLUDED.failed_at,
                        error = EXCLUDED.error,
                        attempt = EXCLUDED.attempt,
                        replayable = EXCLUDED.replayable
                    """,
                    (
                        record.id,
                        record.agent_name,
                        datetime.now(UTC),
                        record.error,
                        record.attempt,
                    ),
                )
            connection.commit()

    async def submit(
        self,
        agent_name: str,
        work: Callable[[], Awaitable[Response]],
        *,
        idempotency_key: str | None = None,
        max_retries: int = 0,
        replay_of: str | None = None,
    ) -> Response:
        if idempotency_key:
            existing = self._lookup_idempotent(idempotency_key)
            if existing is not None:
                return existing
        task_id = str(uuid4())
        record = TaskRecord(
            id=task_id,
            agent_name=agent_name,
            status="queued",
            created_at=datetime.now(UTC),
            idempotency_key=idempotency_key,
            replay_of=replay_of,
        )
        self._insert_task(record)

        async with self._semaphore:
            attempt = 0
            while True:
                attempt += 1
                started_at = datetime.now(UTC)
                self._update_task(task_id, status="running", started_at=started_at, attempt=attempt)
                try:
                    response = await work()
                except Exception as exc:
                    if attempt <= max_retries:
                        self._update_task(
                            task_id,
                            status="queued",
                            started_at=started_at,
                            finished_at=None,
                            error=str(exc),
                            attempt=attempt,
                        )
                        continue
                    failed_record = TaskRecord(
                        id=task_id,
                        agent_name=agent_name,
                        status="failed",
                        created_at=record.created_at,
                        started_at=started_at,
                        finished_at=datetime.now(UTC),
                        error=str(exc),
                        idempotency_key=idempotency_key,
                        attempt=attempt,
                        replay_of=replay_of,
                    )
                    self._update_task(
                        task_id,
                        status="failed",
                        started_at=started_at,
                        finished_at=failed_record.finished_at,
                        error=str(exc),
                        attempt=attempt,
                    )
                    self._record_dead_letter(failed_record)
                    raise
                response_json = _serialize_response(response)
                self._update_task(
                    task_id,
                    status="completed",
                    started_at=started_at,
                    finished_at=datetime.now(UTC),
                    attempt=attempt,
                    response_json=response_json,
                )
                return response

    async def replay(
        self,
        task_id: str,
        agent_name: str,
        work: Callable[[], Awaitable[Response]],
        *,
        max_retries: int = 0,
    ) -> Response:
        return await self.submit(
            agent_name,
            work,
            max_retries=max_retries,
            replay_of=task_id,
        )

    def stats(self) -> dict[str, int]:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT status, COUNT(*) AS count FROM {self.table_name} GROUP BY status"
                )
                rows = cursor.fetchall()
                cursor.execute(
                    f"SELECT COUNT(*) FROM {self.table_name}_dead_letter"
                )
                dead_letter = cursor.fetchone()
        stats = {status: count for status, count in rows}
        stats["dead_letter"] = int(dead_letter[0]) if dead_letter else 0
        return stats


class LoadBalancer:
    """Replica selection strategies for distributed agents."""

    def __init__(self, strategy: str = "round_robin") -> None:
        if strategy not in {"round_robin", "least_loaded", "random"}:
            raise ValueError(f"Unsupported load balancing strategy: {strategy}")
        self.strategy = strategy
        self._counters: defaultdict[str, int] = defaultdict(int)

    def select(self, agent_name: str, replicas: list[str], inflight: dict[str, int]) -> str:
        if not replicas:
            raise LookupError(f"No replicas registered for agent '{agent_name}'")
        if self.strategy == "random":
            return random.choice(replicas)
        if self.strategy == "least_loaded":
            return min(replicas, key=lambda replica: (inflight.get(replica, 0), replica))
        index = self._counters[agent_name] % len(replicas)
        self._counters[agent_name] += 1
        return replicas[index]


class DistributedRuntime:
    """Runtime that can dispatch logical agents across multiple replicas."""

    def __init__(
        self,
        *,
        load_balancer: LoadBalancer | None = None,
        task_queue: SQLiteTaskQueue | PostgreSQLTaskQueue | None = None,
        session_affinity: bool = False,
    ) -> None:
        self.load_balancer = load_balancer or LoadBalancer()
        if task_queue is not None:
            self.task_queue = task_queue
        else:
            database_url = resolve_database_url(env_vars=("LOGOS_ADK_TASK_QUEUE_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"))
            self.task_queue = (
                PostgreSQLTaskQueue(database_url)
                if database_url
                else SQLiteTaskQueue()
            )
        self.session_affinity = session_affinity
        self._replicas: dict[str, dict[str, Agent]] = defaultdict(dict)
        self._inflight: defaultdict[str, int] = defaultdict(int)
        self._session_affinity: dict[str, dict[str, str]] = defaultdict(dict)

    def register(self, logical_name: str, agent: Agent, *, replica_id: str | None = None) -> str:
        """Register a concrete agent replica under a logical agent name."""
        assigned_replica_id = replica_id or f"{logical_name}-{len(self._replicas[logical_name]) + 1}"
        if assigned_replica_id in self._replicas[logical_name]:
            raise ValueError(
                f"Replica '{assigned_replica_id}' already registered for agent '{logical_name}'"
            )
        self._replicas[logical_name][assigned_replica_id] = agent
        agent._runtime = self
        return assigned_replica_id

    def stop(self, logical_name: str, replica_id: str) -> None:
        """Remove a replica from a logical agent pool."""
        replica = self._replicas.get(logical_name, {}).pop(replica_id, None)
        if replica is not None:
            replica._runtime = None
            self._inflight.pop(replica_id, None)
            if logical_name in self._session_affinity:
                self._session_affinity[logical_name] = {
                    session_id: pinned_replica_id
                    for session_id, pinned_replica_id in self._session_affinity[logical_name].items()
                    if pinned_replica_id != replica_id
                }

    def list_agents(self) -> list[str]:
        return sorted(self._replicas)

    def replicas(self, logical_name: str) -> list[str]:
        return sorted(self._replicas.get(logical_name, {}))

    def get_agent(self, logical_name: str) -> Agent | None:
        replicas = self._replicas.get(logical_name, {})
        if not replicas:
            return None
        first_replica_id = sorted(replicas)[0]
        return replicas[first_replica_id]

    def get_graph(self, logical_name: str) -> dict[str, Any] | None:
        """Get the static graph definition from a logical agent if it is a Graph."""
        agent = self.get_agent(logical_name)
        if agent is None:
            return None
        if hasattr(agent, "get_graph_definition"):
            return agent.get_graph_definition()
        return None

    def _select_replica(self, logical_name: str, *, session_id: str | None = None) -> tuple[str, Agent]:
        replicas = self._replicas.get(logical_name, {})
        if self.session_affinity and session_id is not None:
            pinned_replica_id = self._session_affinity.get(logical_name, {}).get(session_id)
            if pinned_replica_id in replicas:
                return pinned_replica_id, replicas[pinned_replica_id]
        replica_ids = sorted(replicas)
        replica_id = self.load_balancer.select(logical_name, replica_ids, self._inflight)
        if self.session_affinity and session_id is not None:
            self._session_affinity[logical_name][session_id] = replica_id
        return replica_id, replicas[replica_id]

    async def dispatch(
        self,
        logical_name: str,
        prompt: str,
        *,
        session_id: str | None = None,
        idempotency_key: str | None = None,
        max_retries: int = 0,
        **kwargs: Any,
    ) -> Response:
        """Queue and execute a run against a selected replica."""
        replica_id, agent = self._select_replica(logical_name, session_id=session_id)

        async def work() -> Response:
            self._inflight[replica_id] += 1
            try:
                return await agent.run(prompt, session_id=session_id, **kwargs)
            finally:
                self._inflight[replica_id] -= 1

        return await self.task_queue.submit(
            logical_name,
            work,
            idempotency_key=idempotency_key,
            max_retries=max_retries,
        )

    async def replay(
        self,
        task_id: str,
        logical_name: str,
        prompt: str,
        *,
        session_id: str | None = None,
        max_retries: int = 0,
    ) -> Response:
        replica_id, agent = self._select_replica(logical_name, session_id=session_id)

        async def work() -> Response:
            self._inflight[replica_id] += 1
            try:
                return await agent.run(prompt, session_id=session_id)
            finally:
                self._inflight[replica_id] -= 1

        return await self.task_queue.replay(
            task_id,
            logical_name,
            work,
            max_retries=max_retries,
        )

    async def stream(
        self,
        logical_name: str,
        prompt: str,
        *,
        session_id: str | None = None,
        **kwargs: Any,
    ):
        """Stream directly from a selected replica."""
        replica_id, agent = self._select_replica(logical_name, session_id=session_id)
        self._inflight[replica_id] += 1
        try:
            async for chunk in agent.stream(prompt, session_id=session_id, **kwargs):
                yield chunk
        finally:
            self._inflight[replica_id] -= 1

    async def stream_ui(
        self,
        logical_name: str,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ):
        """Stream UI events directly from a selected replica."""
        replica_id, agent = self._select_replica(logical_name, session_id=session_id)
        if not hasattr(agent, "stream_ui_events"):
            raise AttributeError(f"Agent '{logical_name}' (replica {replica_id}) does not support stream_ui_events")

        self._inflight[replica_id] += 1
        try:
            async for event in agent.stream_ui_events(
                prompt,
                session_id=session_id,
                correlation_id=correlation_id,
                **kwargs,
            ):
                yield event
        finally:
            self._inflight[replica_id] -= 1

    def stats(self) -> dict[str, Any]:
        """Return runtime and queue statistics."""
        return {
            "agents": {
                agent_name: {
                    "replicas": len(replicas),
                    "replica_ids": sorted(replicas),
                }
                for agent_name, replicas in sorted(self._replicas.items())
            },
            "inflight": dict(self._inflight),
            "queue": self.task_queue.stats(),
            "load_balancer": self.load_balancer.strategy,
            "session_affinity": {
                "enabled": self.session_affinity,
                "sessions": sum(len(mapping) for mapping in self._session_affinity.values()),
            },
        }
