"""AlbusOS reference HTTP server.

Provides production-ready endpoints for agent execution, webhook
routing, job management, approval workflows, and trace retrieval.

Requires the ``[server]`` optional dependency group::

    pip install albusos[server]

Usage::

    # Programmatic
    from albusos.server import create_app
    app = create_app()

    # CLI
    uvicorn albusos.server:app --host 0.0.0.0 --port 8000
"""

from __future__ import annotations

import logging
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Lazy FastAPI import — fail fast with a clear message
# ---------------------------------------------------------------------------

try:
    from fastapi import FastAPI, HTTPException
    from fastapi.responses import JSONResponse
except ImportError as exc:
    raise ImportError(
        "The AlbusOS server requires FastAPI. Install it with: pip install albusos[server]"
    ) from exc


# =============================================================================
# Request / Response models
# =============================================================================


class RunRequest(BaseModel):
    """Synchronous agent execution request."""

    message: str
    thread_id: str | None = None
    inputs: dict[str, Any] | None = None


class SubmitRequest(BaseModel):
    """Asynchronous agent submission request."""

    message: str
    inputs: dict[str, Any] | None = None


class ApprovalRequest(BaseModel):
    """Approve or reject a pending checkpoint."""

    approved: bool = True
    approved_by: str | None = None
    approver_role: str | None = None


class WebhookPayload(BaseModel):
    """Incoming webhook payload."""

    data: dict[str, Any] = Field(default_factory=dict)


class HealthResponse(BaseModel):
    """Health check response."""

    status: str = "ok"
    version: str = ""
    agents: list[str] = Field(default_factory=list)


# =============================================================================
# Application state (populated at startup)
# =============================================================================


class _AppState:
    """Mutable application state shared across endpoints.

    Optional integrations (env vars):

    - ``ALBUS_OTEL=1``    — export spans to an OTLP backend
    - ``ALBUS_METRICS=1`` — export Prometheus metrics

    Tracing is always on — no flag needed.
    """

    def __init__(self) -> None:
        self.workspace: Any = None
        self.sessions: dict[str, Any] = {}

        import os

        self.otel: bool = os.getenv("ALBUS_OTEL", "").lower() in ("1", "true", "yes")
        self.metrics: bool = os.getenv("ALBUS_METRICS", "").lower() in ("1", "true", "yes")


_state = _AppState()


# =============================================================================
# Lifespan
# =============================================================================


@asynccontextmanager
async def _lifespan(_app: FastAPI) -> AsyncIterator[None]:
    """Load workspace and stdlib at startup."""
    import albusos

    albusos._ensure_stdlib()

    try:
        _state.workspace = albusos.load_workspace(".")
        logger.info(
            "Server started: %d agents, %d pathways, %d skills",
            len(_state.workspace.agents),
            len(_state.workspace.pathways),
            len(_state.workspace.skills),
        )
    except FileNotFoundError:
        logger.info("No workspace found at '.', running with stdlib only")
    except Exception:
        logger.exception("Failed to load workspace")

    yield

    _state.sessions.clear()


# =============================================================================
# App factory
# =============================================================================


def create_app() -> FastAPI:
    """Create and configure the FastAPI application."""
    application = FastAPI(
        title="AlbusOS",
        description="Reference HTTP server for AlbusOS agent execution",
        version=_get_version(),
        lifespan=_lifespan,
    )

    application.include_router(_build_router())
    return application


def _get_version() -> str:
    try:
        from importlib.metadata import version

        return version("albusos")
    except Exception:
        return "0.0.0"


# =============================================================================
# Router
# =============================================================================


def _build_router():
    """Build the API router with all endpoints."""
    from fastapi import APIRouter

    router = APIRouter()

    # -----------------------------------------------------------------
    # Health
    # -----------------------------------------------------------------

    @router.get("/health", response_model=HealthResponse)
    async def health() -> HealthResponse:
        agents = _state.workspace.agents if _state.workspace else []
        return HealthResponse(
            status="ok",
            version=_get_version(),
            agents=agents,
        )

    # -----------------------------------------------------------------
    # POST /agent/{agent_id}/run — synchronous execution
    # -----------------------------------------------------------------

    @router.post("/agent/{agent_id}/run")
    async def agent_run(agent_id: str, req: RunRequest) -> JSONResponse:
        session = await _get_or_create_session(agent_id)
        result = await session.run(
            req.message,
            thread_id=req.thread_id,
            inputs=req.inputs or {},
        )
        payload: dict[str, Any] = {
            "success": result.success,
            "response": result.response,
            "tokens_used": result.tokens_used,
            "steps_taken": result.steps_taken,
        }
        # Trace is always available — include it in the response
        if session.last_trace:
            payload["trace"] = session.last_trace.model_dump()
        return JSONResponse(payload)

    # -----------------------------------------------------------------
    # POST /agent/{agent_id}/submit — async submission
    # -----------------------------------------------------------------

    @router.post("/agent/{agent_id}/submit", status_code=202)
    async def agent_submit(agent_id: str, req: SubmitRequest) -> JSONResponse:
        import albusos

        if _state.workspace is None:
            raise HTTPException(500, "No workspace loaded")
        try:
            defn = _state.workspace.agent(agent_id)
        except KeyError as exc:
            raise HTTPException(404, str(exc)) from None

        message = req.message
        if req.inputs:
            message = req.inputs.get("message", message) or message  # type: ignore[assignment]
        job = await albusos.submit(defn, message)
        return JSONResponse(
            {"job_id": job.job_id, "status": job.status},
            status_code=202,
        )

    # -----------------------------------------------------------------
    # GET /job/{job_id} — poll job status
    # -----------------------------------------------------------------

    @router.get("/job/{job_id}")
    async def job_status(job_id: str) -> JSONResponse:
        from core.workflow import get_job

        job = get_job(job_id)
        if job is None:
            raise HTTPException(404, f"Job {job_id} not found")
        payload: dict[str, Any] = {
            "job_id": job.job_id,
            "status": job.status,
            "submitted_at": job.submitted_at,
            "completed_at": job.completed_at,
            "error": job.error,
        }
        if job.result is not None:
            payload["result"] = {
                "success": job.result.success,
                "response": job.result.response,
                "tokens_used": job.result.tokens_used,
            }
        return JSONResponse(payload)

    # -----------------------------------------------------------------
    # POST /approval/{execution_id} — approve a pending checkpoint
    # -----------------------------------------------------------------

    @router.post("/approval/{execution_id}")
    async def approve_checkpoint(
        execution_id: str,
        req: ApprovalRequest,
    ) -> JSONResponse:
        if not req.approved:
            return JSONResponse(
                {
                    "approved": False,
                    "execution_id": execution_id,
                    "message": "Checkpoint rejected",
                }
            )

        session = _find_session_for_execution(execution_id)
        if session is None:
            raise HTTPException(
                404,
                f"No active session found for execution {execution_id}",
            )

        try:
            result = await session.resume_from_checkpoint(execution_id)
            return JSONResponse(
                {
                    "approved": True,
                    "execution_id": execution_id,
                    "approved_by": req.approved_by,
                    "success": result.success,
                    "response": result.response,
                }
            )
        except ValueError as e:
            raise HTTPException(400, str(e)) from e

    # -----------------------------------------------------------------
    # GET /trace/{execution_id} — retrieve execution trace
    # -----------------------------------------------------------------

    @router.get("/trace/{execution_id}")
    async def get_trace(execution_id: str) -> JSONResponse:
        # Try each session's in-memory traces
        for session in _state.sessions.values():
            trace = session.get_trace(execution_id)
            if trace is not None:
                return JSONResponse(trace.model_dump())

        # Fallback: try the default store directly
        from core.state import FileStateStore

        store = FileStateStore()
        trace = await store.load_trace(execution_id)
        if trace is not None:
            return JSONResponse(trace.model_dump())

        raise HTTPException(404, f"Trace {execution_id} not found")

    # -----------------------------------------------------------------
    # POST /webhook/{topic} — webhook ingestion
    # -----------------------------------------------------------------

    @router.post("/webhook/{topic}", status_code=202)
    async def webhook_ingest(topic: str, payload: WebhookPayload) -> JSONResponse:
        from core.pathways.streaming.event_bus import WEBHOOK_BUS

        delivered = WEBHOOK_BUS.publish(topic, payload.data)
        return JSONResponse(
            {"topic": topic, "delivered": delivered},
            status_code=202,
        )

    return router


# =============================================================================
# Helpers
# =============================================================================


async def _get_or_create_session(agent_id: str) -> Any:
    """Get or create a Session for the given agent."""
    if agent_id in _state.sessions:
        return _state.sessions[agent_id]

    if _state.workspace is None:
        raise HTTPException(500, "No workspace loaded")

    try:
        session = _state.workspace.session(
            agent_id,
            otel=_state.otel,
            metrics=_state.metrics,
        )
    except KeyError as exc:
        raise HTTPException(404, str(exc)) from None

    _state.sessions[agent_id] = session
    return session


def _find_session_for_execution(execution_id: str) -> Any | None:
    """Find the session that has a checkpoint for this execution."""
    for session in _state.sessions.values():
        record = session._vm._executions.get(execution_id)
        if record is not None:
            return session
    return None


# =============================================================================
# Module-level app for `uvicorn albusos.server:app`
# =============================================================================

app = create_app()

__all__ = [
    "app",
    "create_app",
]
