"""Web server for Logos ADK agents."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
import json
import os
from typing import Any
from collections import defaultdict
import time
from uuid import uuid4

from fastapi import Request

from logos_adk.observe import get_native_exporter, get_observability_store, Metric
from logos_adk.config import RateLimitConfig
from logos_adk.security import RateLimiter, build_rate_limiter
from logos_adk.serve_security import (
    WORKSPACE_HEADER,
    load_api_keys_config,
    required_scope_for_request,
    resolve_workspace_id,
)
from logos_adk.version import __version__


class Metrics:
    """Simple in-memory metrics collector."""

    def __init__(self):
        self.total_requests = 0
        self.total_tokens_in = 0
        self.total_tokens_out = 0
        self.agent_requests: dict[str, int] = defaultdict(int)
        self.agent_tokens: dict[str, dict] = defaultdict(lambda: {"in": 0, "out": 0})
        self.start_time = time.time()

    def record(self, agent_name: str, tokens_in: int, tokens_out: int):
        self.total_requests += 1
        self.total_tokens_in += tokens_in
        self.total_tokens_out += tokens_out
        self.agent_requests[agent_name] += 1
        self.agent_tokens[agent_name]["in"] += tokens_in
        self.agent_tokens[agent_name]["out"] += tokens_out

    def get_stats(self) -> dict:
        uptime = time.time() - self.start_time
        return {
            "total_requests": self.total_requests,
            "total_tokens": self.total_tokens_in + self.total_tokens_out,
            "tokens_in": self.total_tokens_in,
            "tokens_out": self.total_tokens_out,
            "uptime_seconds": round(uptime, 2),
            "by_agent": {
                name: {
                    "requests": self.agent_requests[name],
                    "tokens_in": self.agent_tokens[name]["in"],
                    "tokens_out": self.agent_tokens[name]["out"],
                }
                for name in self.agent_requests
            },
        }


_global_metrics = Metrics()


def get_metrics() -> Metrics:
    """Get the global metrics instance."""
    return _global_metrics


def _metrics_agent_name_for_path(path: str, available_agents: list[str]) -> str | None:
    """Infer logical agent name for metrics from a request path."""
    normalized = path.strip("/") or ""
    parts = normalized.split("/") if normalized else []
    execution_routes = {"run", "stream", "batch", "stream-ui", "resume"}
    if not parts:
        return None
    if parts[0] in execution_routes and len(available_agents) == 1:
        return available_agents[0]
    if len(parts) >= 2 and parts[1] in execution_routes and parts[0] in available_agents:
        return parts[0]
    return None


def _ensure_native_observer(agent: Any) -> None:
    exporter = get_native_exporter()
    observers = getattr(getattr(agent, "_config", None), "observers", None)
    if isinstance(observers, list) and exporter not in observers:
        observers.insert(0, exporter)


def _resolve_rate_limit_config(
    agents: list[Any],
    *,
    rate_limit_enabled: bool | None,
    rate_limit_per_minute: int | None,
    rate_limit_burst: int | None,
    rate_limit_backend: str | None,
    rate_limit_dsn: str | None,
) -> RateLimitConfig:
    for agent in agents:
        config = getattr(getattr(agent, "_config", None), "security", None)
        if config is not None:
            base = config.rate_limit
            break
    else:
        base = RateLimitConfig()
    return RateLimitConfig(
        enabled=base.enabled if rate_limit_enabled is None else rate_limit_enabled,
        backend=rate_limit_backend or base.backend,
        dsn=rate_limit_dsn or base.dsn,
        namespace=base.namespace,
        table_name=base.table_name,
        limit_per_minute=rate_limit_per_minute or base.limit_per_minute,
        burst=rate_limit_burst if rate_limit_burst is not None else base.burst,
        window_seconds=base.window_seconds,
    )


def create_app(
    agents: list[Any],
    *,
    runtime: Any | None = None,
    learning_store: Any | None = None,
    prompt_registry: Any | None = None,
    cors: bool = True,
    api_key: str | None = None,
    api_keys: dict[str, dict[str, Any]] | None = None,
    require_workspace: bool | None = None,
    timeout: float = 120,
    max_body_size: int = 1_048_576,
    rate_limit_enabled: bool | None = None,
    rate_limit_per_minute: int | None = None,
    rate_limit_burst: int | None = None,
    rate_limit_backend: str | None = None,
    rate_limit_dsn: str | None = None,
    chronos_autostart: bool = True,
    chronos_poll_interval_seconds: float = 1.0,
) -> Any:
    """Create a FastAPI application serving the given agents.

    Args:
        agents: List of Agent instances to serve.
        runtime: Optional distributed runtime used to dispatch requests.
        cors: Enable CORS with allow_origins=["*"]. Default True.
        api_key: Optional legacy API key for Bearer auth on all endpoints except /health.
        timeout: Per-request timeout in seconds. Default 120.
        max_body_size: Max request body size in bytes. Default 1MB.

    Returns:
        A FastAPI application instance.
    """
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse

    from logos_adk.learning import get_learning_store, get_prompt_registry
    from logos_adk.chronos import Chronos
    from logos_adk.serve_routes import create_agent_routes

    rate_limit_config = _resolve_rate_limit_config(
        agents,
        rate_limit_enabled=rate_limit_enabled,
        rate_limit_per_minute=rate_limit_per_minute,
        rate_limit_burst=rate_limit_burst,
        rate_limit_backend=rate_limit_backend,
        rate_limit_dsn=rate_limit_dsn,
    )
    rate_limiter = (
        build_rate_limiter(rate_limit_config) if rate_limit_config.enabled else RateLimiter()
    )

    continuous_eval_runners: list[Any] = []
    chronos_runners: list[Any] = []

    @asynccontextmanager
    async def lifespan(app: Any):
        for runner in continuous_eval_runners:
            restore = getattr(runner, "restore_jobs", None)
            if callable(restore):
                try:
                    restore()
                except Exception:
                    continue
            start = getattr(runner, "start", None)
            if callable(start):
                try:
                    await start()
                except Exception:
                    continue
        if chronos_autostart:
            for runner in chronos_runners:
                start = getattr(runner, "start", None)
                if callable(start):
                    try:
                        result = start()
                        if asyncio.iscoroutine(result):
                            await result
                    except Exception:
                        continue
        yield
        for runner in chronos_runners:
            stop = getattr(runner, "stop", None)
            if callable(stop):
                try:
                    result = stop()
                    if asyncio.iscoroutine(result):
                        await result
                except Exception:
                    continue
        for runner in continuous_eval_runners:
            stop = getattr(runner, "stop", None)
            if callable(stop):
                try:
                    await stop()
                except Exception:
                    continue
        await rate_limiter.close()
        for agent in agents:
            for exporter in getattr(getattr(agent, "_config", None), "observers", []) or []:
                flush = getattr(exporter, "flush", None)
                if flush is None:
                    continue
                try:
                    result = flush()
                    if asyncio.iscoroutine(result):
                        await result
                except Exception:
                    continue

    app = FastAPI(title="Logos ADK", version=__version__, lifespan=lifespan)

    # Custom exception handler for consistent error format
    from fastapi import HTTPException
    from fastapi.exceptions import RequestValidationError
    @app.exception_handler(HTTPException)
    async def http_exception_handler(request: Request, exc: HTTPException):
        """Handle HTTP exceptions with consistent error format."""
        error_type = exc.detail if isinstance(exc.detail, str) else "error"
        if error_type == "guardrail":
            error_type = "guardrail"
        elif error_type == "timeout":
            error_type = "timeout"
        elif "internal" in error_type.lower() or "Agent execution failed" in error_type:
            error_type = "internal"
        return JSONResponse(
            status_code=exc.status_code,
            content={"error": error_type, "message": exc.detail},
        )

    @app.exception_handler(RequestValidationError)
    async def validation_exception_handler(request: Request, exc: RequestValidationError):
        """Handle validation errors with consistent error format."""
        return JSONResponse(
            status_code=422,
            content={"error": "validation", "message": str(exc.errors())},
        )

    # Mount artifacts routes (global, not per-agent)
    from logos_adk.serve_routes.routes.artifacts import router as artifacts_router
    from logos_adk.serve_routes.routes.quality import router as quality_router

    app.include_router(artifacts_router)
    app.include_router(quality_router)

    # Campaign routes will be added when teams are registered

    @app.exception_handler(Exception)
    async def general_exception_handler(request: Request, exc: Exception):
        """Handle general exceptions with consistent error format."""
        return JSONResponse(
            status_code=500,
            content={"error": "internal_error", "message": str(exc)},
        )

    app.state.continuous_eval_runners = continuous_eval_runners
    app.state.chronos_runners = chronos_runners
    observability = get_observability_store()
    api_key_registry = load_api_keys_config(api_key=api_key, api_keys=api_keys)
    resolved_learning_store = learning_store or get_learning_store()
    resolved_prompt_registry = prompt_registry or get_prompt_registry()
    if require_workspace is None:
        require_workspace = os.getenv("LOGOS_ADK_REQUIRE_WORKSPACE", "").lower() in {
            "1",
            "true",
            "yes",
        }
    if api_key_registry and not require_workspace:
        require_workspace = os.getenv("LOGOS_ADK_AUTH_REQUIRE_WORKSPACE", "").lower() in {
            "1",
            "true",
            "yes",
        }

    for agent in agents:
        _ensure_native_observer(agent)
    if runtime is not None and hasattr(runtime, "list_agents"):
        for logical_name in runtime.list_agents():
            agent = runtime.get_agent(logical_name)
            if agent is not None:
                _ensure_native_observer(agent)

    # CORS
    if cors:
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )

    # Auth and permission middleware
    if api_key_registry:

        @app.middleware("http")
        async def auth_middleware(request: Request, call_next):
            if request.url.path.startswith("/health"):
                return await call_next(request)
            auth_header = request.headers.get("Authorization", "")
            query_api_key = request.query_params.get("api_key") or request.query_params.get("token")
            token = None
            if auth_header.startswith("Bearer "):
                token = auth_header.split(" ", 1)[1].strip()
            elif query_api_key:
                token = query_api_key
            principal = api_key_registry.get(token or "")
            if principal is None:
                return JSONResponse(
                    status_code=401,
                    content={"error": "unauthorized", "message": "Invalid or missing API key"},
                )
            workspace_id = resolve_workspace_id(request.headers, request.query_params)
            required_scope = required_scope_for_request(request.url.path, request.method)
            if required_scope is not None and not principal.has_scope(required_scope):
                return JSONResponse(
                    status_code=403,
                    content={
                        "error": "forbidden",
                        "message": f"Missing required scope: {required_scope}",
                    },
                )
            if require_workspace and workspace_id is None:
                return JSONResponse(
                    status_code=400,
                    content={
                        "error": "workspace_required",
                        "message": f"Missing {WORKSPACE_HEADER} header",
                    },
                )
            if workspace_id is not None and not principal.allows_workspace(workspace_id):
                return JSONResponse(
                    status_code=403,
                    content={"error": "workspace_forbidden", "message": "Workspace access denied"},
                )
            request.state.principal = principal
            request.state.workspace_id = workspace_id
            return await call_next(request)

    else:

        @app.middleware("http")
        async def workspace_context_middleware(request: Request, call_next):
            request.state.workspace_id = resolve_workspace_id(request.headers, request.query_params)
            return await call_next(request)

    # Max body size middleware
    @app.middleware("http")
    async def body_size_middleware(request: Request, call_next):
        content_length = request.headers.get("content-length")
        if content_length and int(content_length) > max_body_size:
            observability.log(
                level="warning",
                message="Rejected oversized request body",
                source="http",
                request_id=request.headers.get("x-request-id"),
                attributes={"path": request.url.path, "content_length": int(content_length)},
            )
            return JSONResponse(
                status_code=413,
                content={"error": "payload_too_large", "message": "Request body too large"},
            )
        return await call_next(request)

    @app.middleware("http")
    async def rate_limit_middleware(request: Request, call_next):
        if request.url.path.startswith("/health") or not rate_limit_config.enabled:
            return await call_next(request)
        client_host = request.client.host if request.client else "unknown"
        key = f"{client_host}:{request.url.path}"
        allowed, retry_after = await rate_limiter.allow(key)
        if not allowed:
            observability.log(
                level="warning",
                message="Rate limit exceeded",
                source="security",
                request_id=request.headers.get("x-request-id"),
                attributes={
                    "path": request.url.path,
                    "client": client_host,
                    "retry_after": retry_after,
                },
            )
            return JSONResponse(
                status_code=429,
                content={
                    "error": "rate_limited",
                    "message": "Rate limit exceeded",
                    "retry_after": retry_after,
                },
                headers={"Retry-After": str(int(retry_after) or 1)},
            )
        return await call_next(request)

    @app.middleware("http")
    async def observability_middleware(request: Request, call_next):
        started_at = time.time()
        trace_id = request.headers.get("x-trace-id") or uuid4().hex
        request_id = request.headers.get("x-request-id") or uuid4().hex[:12]
        request.state.trace_id = trace_id
        request.state.request_id = request_id
        request.state.started_at = started_at

        observability.log(
            level="info",
            message="Request started",
            source="http",
            trace_id=trace_id,
            request_id=request_id,
            attributes={
                "path": request.url.path,
                "method": request.method,
            },
        )
        try:
            response = await call_next(request)
        except Exception as exc:
            duration_ms = (time.time() - started_at) * 1000
            observability.record_metric(
                Metric(
                    name="adk_request_latency_ms",
                    value=duration_ms,
                    unit="ms",
                    attributes={
                        "path": request.url.path,
                        "method": request.method,
                        "status": "500",
                    },
                )
            )
            observability.record_metric(
                Metric(
                    name="adk_error_total",
                    value=1.0,
                    unit="count",
                    attributes={"path": request.url.path, "method": request.method},
                )
            )
            observability.log(
                level="error",
                message="Unhandled server error",
                source="http",
                trace_id=trace_id,
                request_id=request_id,
                attributes={"path": request.url.path, "error": str(exc)},
            )
            raise

        duration_ms = (time.time() - started_at) * 1000
        status = getattr(response, "status_code", 200)
        observability.record_metric(
            Metric(
                name="adk_http_requests_total",
                value=1.0,
                unit="count",
                attributes={
                    "path": request.url.path,
                    "method": request.method,
                    "status": str(status),
                },
            )
        )
        observability.record_metric(
            Metric(
                name="adk_request_latency_ms",
                value=duration_ms,
                unit="ms",
                attributes={
                    "path": request.url.path,
                    "method": request.method,
                    "status": str(status),
                },
            )
        )
        if status >= 500:
            observability.record_metric(
                Metric(
                    name="adk_error_total",
                    value=1.0,
                    unit="count",
                    attributes={
                        "path": request.url.path,
                        "method": request.method,
                        "status": str(status),
                    },
                )
            )
        observability.log(
            level="error" if status >= 500 else "warning" if status >= 400 else "info",
            message="Request completed",
            source="http",
            trace_id=trace_id,
            request_id=request_id,
            attributes={
                "path": request.url.path,
                "method": request.method,
                "status": status,
                "duration_ms": round(duration_ms, 2),
            },
        )
        metrics_agent = _metrics_agent_name_for_path(request.url.path, list(agent_names))
        if metrics_agent is not None:
            get_metrics().record(metrics_agent, 0, 0)
        response.headers["X-Trace-Id"] = trace_id
        response.headers["X-Request-Id"] = request_id
        workspace_id = getattr(request.state, "workspace_id", None)
        if workspace_id is not None:
            response.headers[WORKSPACE_HEADER] = workspace_id
        return response

    # Health endpoint
    if runtime is not None and hasattr(runtime, "list_agents"):
        agent_names = runtime.list_agents()
    else:
        agent_names = [a.name for a in agents]

    @app.get("/health")
    async def health():
        payload = {"status": "ok", "agents": agent_names}
        if runtime is not None and hasattr(runtime, "stats"):
            runtime_stats = runtime.stats()
            payload["runtime"] = runtime_stats
            queue_stats = runtime_stats.get("queue", {})
            backlog = float(queue_stats.get("queued", 0) + queue_stats.get("running", 0))
            observability.record_metric(
                Metric(
                    name="adk_queue_backlog",
                    value=backlog,
                    unit="count",
                    attributes={"source": "runtime"},
                )
            )
        if agents:
            payload["providers"] = [agent.health() for agent in agents if hasattr(agent, "health")]
        return payload

    @app.get("/health/live")
    async def health_live():
        return {"status": "alive"}

    @app.get("/health/ready")
    async def health_ready():
        provider_states = []
        for agent in agents:
            if hasattr(agent, "health"):
                provider_states.extend(agent.health().get("providers", []))
        ready = (
            all(state.get("healthy", True) for state in provider_states)
            if provider_states
            else True
        )
        return JSONResponse(
            status_code=200 if ready else 503,
            content={"status": "ready" if ready else "degraded", "providers": provider_states},
        )

    @app.get("/metrics")
    async def metrics():
        return get_metrics().get_stats()

    @app.get("/metrics/prometheus")
    async def metrics_prometheus():
        return PlainTextResponse(observability.prometheus_text())

    @app.get("/observability/summary")
    async def observability_summary():
        return observability.summary()

    @app.get("/observability/spans")
    async def observability_spans(
        limit: int = 100,
        trace_id: str | None = None,
        agent_name: str | None = None,
        status: str | None = None,
        tenant_id: str | None = None,
        tool_name: str | None = None,
    ):
        attributes_filter = {
            key: value
            for key, value in {
                "tenant_id": tenant_id,
                "tool_name": tool_name,
            }.items()
            if value is not None
        }
        return {
            "items": observability.query_spans(
                limit=limit,
                trace_id=trace_id,
                agent_name=agent_name,
                status=status,
                attributes_filter=attributes_filter or None,
            )
        }

    @app.get("/observability/metrics")
    async def observability_metrics(
        limit: int = 200,
        name: str | None = None,
        tenant_id: str | None = None,
        tool_name: str | None = None,
    ):
        attributes_filter = {
            key: value
            for key, value in {
                "tenant_id": tenant_id,
                "tool_name": tool_name,
            }.items()
            if value is not None
        }
        return {
            "items": observability.query_metrics(
                limit=limit,
                name=name,
                attributes_filter=attributes_filter or None,
            )
        }

    @app.get("/observability/logs")
    async def observability_logs(
        limit: int = 200,
        level: str | None = None,
        trace_id: str | None = None,
        source: str | None = None,
        tenant_id: str | None = None,
        tool_name: str | None = None,
    ):
        attributes_filter = {
            key: value
            for key, value in {
                "tenant_id": tenant_id,
                "tool_name": tool_name,
            }.items()
            if value is not None
        }
        return {
            "items": observability.query_logs(
                limit=limit,
                level=level,
                trace_id=trace_id,
                source=source,
                attributes_filter=attributes_filter or None,
            )
        }

    @app.get("/observability/alerts")
    async def observability_alerts(
        limit: int = 100,
        severity: str | None = None,
        rule_name: str | None = None,
        tenant_id: str | None = None,
        tool_name: str | None = None,
    ):
        attributes_filter = {
            key: value
            for key, value in {
                "tenant_id": tenant_id,
                "tool_name": tool_name,
            }.items()
            if value is not None
        }
        return {
            "items": observability.query_alerts(
                limit=limit,
                severity=severity,
                rule_name=rule_name,
                attributes_filter=attributes_filter or None,
            )
        }

    @app.get("/observability/grafana/dashboard")
    async def observability_grafana_dashboard():
        return observability.grafana_dashboard()

    def _resolve_named_agent(agent_name: str | None) -> Any | None:
        if not agent_name:
            return None
        if runtime is not None and hasattr(runtime, "get_agent"):
            return runtime.get_agent(agent_name)
        for agent in agents:
            if agent.name == agent_name:
                return agent
        return None

    def _observability_stream_payload(
        *,
        selected_agent: str | None,
        session_id: str | None,
        trace_id: str | None,
        status: str | None,
        level: str | None,
        tenant_id: str | None,
        tool_name: str | None,
    ) -> dict[str, Any]:
        attributes_filter = {
            key: value
            for key, value in {
                "tenant_id": tenant_id,
                "tool_name": tool_name,
            }.items()
            if value is not None
        }
        payload: dict[str, Any] = {
            "summary": observability.summary(),
            "traces": observability.query_spans(
                limit=100,
                trace_id=trace_id,
                agent_name=selected_agent,
                status=status,
                attributes_filter=attributes_filter or None,
            ),
            "logs": observability.query_logs(
                limit=200,
                level=level,
                trace_id=trace_id,
                attributes_filter=attributes_filter or None,
            ),
            "alerts": observability.query_alerts(
                limit=100,
                severity=level,
                attributes_filter=attributes_filter or None,
            ),
        }
        agent = _resolve_named_agent(selected_agent)
        if agent is not None:
            payload["state_versions"] = agent.list_state_versions(session_id, limit=50)
            payload["state_backend"] = agent.state_store_info()
            payload["state_export"] = agent.export_state(session_id)
        else:
            payload["state_versions"] = []
            payload["state_backend"] = None
            payload["state_export"] = None
        return payload

    @app.get("/observability/stream")
    async def observability_stream(
        request: Request,
        selected_agent: str | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        status: str | None = None,
        level: str | None = None,
        tenant_id: str | None = None,
        tool_name: str | None = None,
        interval_ms: int = 1000,
        once: bool = False,
    ):
        interval_seconds = max(interval_ms, 250) / 1000

        async def event_generator():
            last_payload = ""
            while True:
                if await request.is_disconnected():
                    break
                payload = _observability_stream_payload(
                    selected_agent=selected_agent,
                    session_id=session_id,
                    trace_id=trace_id,
                    status=status,
                    level=level,
                    tenant_id=tenant_id,
                    tool_name=tool_name,
                )
                serialized = json.dumps(payload, separators=(",", ":"), default=str)
                if serialized != last_payload or once:
                    yield f"event: snapshot\ndata: {serialized}\n\n"
                    last_payload = serialized
                else:
                    yield "event: keepalive\ndata: {}\n\n"
                if once:
                    break
                await asyncio.sleep(interval_seconds)

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    # Mount agent routes
    registered_names = set()

    class _RuntimeRunnableProxy:
        """Present a logical runtime pool as a single runnable for HTTP serving."""

        def __init__(self, logical_name: str, runtime_obj: Any) -> None:
            self.name = logical_name
            self._runtime = runtime_obj
            self._backing_agent = runtime_obj.get_agent(logical_name)
            self._config = getattr(self._backing_agent, "_config", None)
            self._state_store = getattr(self._backing_agent, "_state_store", None)

        async def run(self, prompt: Any, *, session_id: str | None = None, **kwargs: Any) -> Any:
            return await self._runtime.dispatch(
                self.name,
                prompt,
                session_id=session_id,
                **kwargs,
            )

        async def stream(self, prompt: Any, *, session_id: str | None = None, **kwargs: Any):
            async for chunk in self._runtime.stream(
                self.name,
                prompt,
                session_id=session_id,
                **kwargs,
            ):
                yield chunk

        async def stream_ui_events(
            self,
            prompt: Any,
            *,
            session_id: str | None = None,
            correlation_id: str | None = None,
            **kwargs: Any,
        ):
            async for event in self._runtime.stream_ui(
                self.name,
                prompt,
                session_id=session_id,
                correlation_id=correlation_id,
                **kwargs,
            ):
                yield event

        def get_graph_definition(self) -> dict[str, Any] | None:
            return self._runtime.get_graph(self.name)

    def register_agent(agent_obj: Any, name: str | None = None, is_primary: bool = False):
        logical_name = name or getattr(agent_obj, "name", "unknown")
        if logical_name in registered_names:
            return
        route_base = f"/{logical_name}"
        _ensure_native_observer(agent_obj)

        from logos_adk.serve_routes.routes.state import create_state_routes

        def get_state_store():
            return getattr(agent_obj, "_state_store", None)

        def get_agent_name():
            return getattr(agent_obj, "name", "default")

        state_router = create_state_routes(get_state_store, get_agent_name)

        # Add campaign routes for Teams
        from logos_adk.team import Team
        from logos_adk.serve_routes.routes.campaigns import create_campaign_routes

        if isinstance(agent_obj, Team):
            campaign_router = create_campaign_routes(lambda: agent_obj)
            app.include_router(campaign_router)
            chronos_runners.append(
                Chronos(agent_obj, poll_interval_seconds=chronos_poll_interval_seconds)
            )

        if is_primary:
            primary_router = create_agent_routes(
                agent_obj,
                timeout=timeout,
                runtime=runtime,
                learning_store=resolved_learning_store,
                prompt_registry=resolved_prompt_registry,
                chronos_poll_interval_seconds=chronos_poll_interval_seconds,
                public_name=logical_name,
                route_base="",
            )
            app.include_router(primary_router)
            # Also add state routes at root for primary agent
            app.include_router(state_router)

        router = create_agent_routes(
            agent_obj,
            timeout=timeout,
            runtime=runtime,
            learning_store=resolved_learning_store,
            prompt_registry=resolved_prompt_registry,
            chronos_poll_interval_seconds=chronos_poll_interval_seconds,
            public_name=logical_name,
            route_base=route_base,
        )

        app.include_router(router, prefix=route_base)

        # Add state routes for this agent (with agent prefix)
        state_router = create_state_routes(get_state_store)
        app.include_router(state_router, prefix=route_base)

        registered_names.add(logical_name)

        runner = getattr(router, "_continuous_eval_runner", None)
        if runner is not None:
            continuous_eval_runners.append(runner)
        chronos_runner = getattr(router, "_chronos_runner", None)
        if chronos_runner is not None:
            chronos_runners.append(chronos_runner)

    @app.post("/refresh_agents")
    async def refresh_agents():
        from logos_adk.tools.factory import dynamic_agents

        added = 0
        for name, obj in dynamic_agents.items():
            if name not in registered_names:
                register_agent(obj, name=name)
                added += 1
        return {"status": "ok", "added": added, "total": len(registered_names)}

    if runtime is not None and hasattr(runtime, "list_agents"):
        for logical_name in runtime.list_agents():
            proxy = _RuntimeRunnableProxy(logical_name, runtime)
            if getattr(proxy, "_backing_agent", None) is None:
                continue
            register_agent(proxy, name=logical_name, is_primary=not registered_names)
    for agent in agents:
        register_agent(agent, is_primary=not registered_names)

    return app


def create_app_from_config_paths(
    config_paths: list[str],
    *,
    runtime: Any | None = None,
    learning_store: Any | None = None,
    prompt_registry: Any | None = None,
    cors: bool = True,
    api_key: str | None = None,
    api_keys: dict[str, dict[str, Any]] | None = None,
    require_workspace: bool | None = None,
    timeout: float = 120,
    max_body_size: int = 1_048_576,
    rate_limit_enabled: bool | None = None,
    rate_limit_per_minute: int | None = None,
    rate_limit_burst: int | None = None,
    rate_limit_backend: str | None = None,
    rate_limit_dsn: str | None = None,
    chronos_autostart: bool = True,
    chronos_poll_interval_seconds: float = 1.0,
) -> Any:
    """Create an app from a list of YAML config paths."""
    from logos_adk.runnable_config_loader import load_runnable_from_config

    runnables = [load_runnable_from_config(path) for path in config_paths]
    return create_app(
        runnables,
        runtime=runtime,
        learning_store=learning_store,
        prompt_registry=prompt_registry,
        cors=cors,
        api_key=api_key,
        api_keys=api_keys,
        require_workspace=require_workspace,
        timeout=timeout,
        max_body_size=max_body_size,
        rate_limit_enabled=rate_limit_enabled,
        rate_limit_per_minute=rate_limit_per_minute,
        rate_limit_burst=rate_limit_burst,
        rate_limit_backend=rate_limit_backend,
        rate_limit_dsn=rate_limit_dsn,
        chronos_autostart=chronos_autostart,
        chronos_poll_interval_seconds=chronos_poll_interval_seconds,
    )


def serve(
    agents: list[Any] | None = None,
    *,
    runtime: Any | None = None,
    learning_store: Any | None = None,
    prompt_registry: Any | None = None,
    port: int = 8000,
    host: str = "0.0.0.0",
    cors: bool = True,
    api_key: str | None = None,
    api_keys: dict[str, dict[str, Any]] | None = None,
    require_workspace: bool | None = None,
    timeout: float = 120,
    max_body_size: int = 1_048_576,
    rate_limit_enabled: bool | None = None,
    rate_limit_per_minute: int | None = None,
    rate_limit_burst: int | None = None,
    rate_limit_backend: str | None = None,
    rate_limit_dsn: str | None = None,
    chronos_autostart: bool = True,
    chronos_poll_interval_seconds: float = 1.0,
) -> None:
    """Create and run a web server for the given agents.

    Args:
        agents: List of Agent instances to serve.
        runtime: Optional distributed runtime used to dispatch requests.
        port: Port to listen on. Default 8000.
        host: Host to bind to. Default "0.0.0.0".
        cors: Enable CORS. Default True.
        api_key: Optional legacy API key for Bearer auth.
        timeout: Per-request timeout in seconds. Default 120.
        max_body_size: Max request body in bytes. Default 1MB.
    """
    import uvicorn

    if agents is None:
        agents = []
    if not agents and runtime is None:
        raise ValueError("agents or runtime must be provided")

    app = create_app(
        agents,
        runtime=runtime,
        learning_store=learning_store,
        prompt_registry=prompt_registry,
        cors=cors,
        api_key=api_key,
        api_keys=api_keys,
        require_workspace=require_workspace,
        timeout=timeout,
        max_body_size=max_body_size,
        rate_limit_enabled=rate_limit_enabled,
        rate_limit_per_minute=rate_limit_per_minute,
        rate_limit_burst=rate_limit_burst,
        rate_limit_backend=rate_limit_backend,
        rate_limit_dsn=rate_limit_dsn,
        chronos_autostart=chronos_autostart,
        chronos_poll_interval_seconds=chronos_poll_interval_seconds,
    )
    uvicorn.run(app, host=host, port=port)


def serve_from_config_paths(
    config_paths: list[str],
    *,
    runtime: Any | None = None,
    learning_store: Any | None = None,
    prompt_registry: Any | None = None,
    port: int = 8000,
    host: str = "0.0.0.0",
    cors: bool = True,
    api_key: str | None = None,
    api_keys: dict[str, dict[str, Any]] | None = None,
    require_workspace: bool | None = None,
    timeout: float = 120,
    max_body_size: int = 1_048_576,
    rate_limit_per_minute: int | None = None,
    rate_limit_burst: int | None = None,
    rate_limit_backend: str | None = None,
    rate_limit_dsn: str | None = None,
    chronos_autostart: bool = True,
    chronos_poll_interval_seconds: float = 1.0,
) -> None:
    """Load runnables from YAML config files and serve them."""
    import uvicorn

    app = create_app_from_config_paths(
        config_paths,
        runtime=runtime,
        learning_store=learning_store,
        prompt_registry=prompt_registry,
        cors=cors,
        api_key=api_key,
        api_keys=api_keys,
        require_workspace=require_workspace,
        timeout=timeout,
        max_body_size=max_body_size,
        rate_limit_per_minute=rate_limit_per_minute,
        rate_limit_burst=rate_limit_burst,
        rate_limit_backend=rate_limit_backend,
        rate_limit_dsn=rate_limit_dsn,
        chronos_autostart=chronos_autostart,
        chronos_poll_interval_seconds=chronos_poll_interval_seconds,
    )
    uvicorn.run(app, host=host, port=port)
