"""FastAPI application for research workbench backend.

Provides auth, AG-UI (CopilotKit) endpoint, session proxy to amplifierd,
artifact API, MCP App serving, and static frontend.
"""

import json
import os
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncGenerator

import httpx
from fastapi import Cookie, Depends, FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from agui_adapter import AmplifierdEventTranslator
from artifacts import get_artifact, list_artifacts
from auth import (
    create_session_token,
    invalidate_session_token,
    validate_session_token,
    verify_password,
)

# ---------------------------------------------------------------------------
# Module constants
# ---------------------------------------------------------------------------

AMPLIFIERD_URL: str = os.environ.get("AMPLIFIERD_URL", "http://localhost:8410")
AUTH_USER: str = os.environ.get("AUTH_USER", "admin@localhost")
BUNDLE_NAME: str = os.environ.get("BUNDLE_NAME", "research-workbench")
APPS_DIR: Path = Path(__file__).parent.parent / "bundle" / "apps"

# ---------------------------------------------------------------------------
# Lifespan — shared HTTP client
# ---------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """Create and teardown the shared httpx.AsyncClient for amplifierd."""
    app.state.http_client = httpx.AsyncClient(
        base_url=AMPLIFIERD_URL,
        timeout=30.0,
    )
    yield
    await app.state.http_client.aclose()


# ---------------------------------------------------------------------------
# Application
# ---------------------------------------------------------------------------

app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
)

# ---------------------------------------------------------------------------
# Auth helper
# ---------------------------------------------------------------------------


def _require_auth(session: str | None = Cookie(None)) -> str:
    """Validate the session cookie and return the email, or raise HTTP 401."""
    email = validate_session_token(session or "")
    if not email:
        raise HTTPException(status_code=401, detail="Not authenticated")
    return email


# ---------------------------------------------------------------------------
# SSE helpers
# ---------------------------------------------------------------------------


def _sse(event: Any) -> str:
    """Format an AG-UI event object as an SSE data line."""
    return f"data: {event.model_dump_json()}\n\n"


def _parse_sse(line: str, state: dict[str, Any]) -> dict[str, Any] | None:
    """Parse one SSE line, accumulate state, return event dict when data arrives.

    Args:
        line: A single line from an SSE stream.
        state: Mutable dict holding the current event type (reused across calls).

    Returns:
        A dict ``{event_type, data}`` when a data line is parsed, else ``None``.
    """
    if line.startswith("event:"):
        state["event_type"] = line[len("event:") :].strip()
        return None
    if line.startswith("data:"):
        raw = line[len("data:") :].strip()
        try:
            data = json.loads(raw)
        except (json.JSONDecodeError, ValueError):
            data = raw
        return {"event_type": state.get("event_type", ""), "data": data}
    return None


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


@app.get("/api/health")
async def health() -> dict[str, str]:
    """Return service health status."""
    return {"status": "ok"}


# ---------------------------------------------------------------------------
# Auth endpoints
# ---------------------------------------------------------------------------


@app.post("/api/auth/login")
async def login(request: Request) -> JSONResponse:
    """Validate credentials, create session token, set session cookie."""
    body = await request.json()
    email: str = body.get("email", "")
    password: str = body.get("password", "")

    if email != AUTH_USER or not verify_password(password):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    token = create_session_token(email)

    response = JSONResponse({"token": token, "email": email})
    response.set_cookie(
        key="session",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=86400 * 7,
    )
    return response


@app.get("/api/auth/me")
async def me(email: str = Depends(_require_auth)) -> dict[str, str]:
    """Return the authenticated user's email."""
    return {"email": email}


@app.post("/api/auth/logout")
async def logout(
    response: Response,
    session: str | None = Cookie(None),
) -> dict[str, str]:
    """Invalidate the session token and clear the session cookie."""
    if session:
        invalidate_session_token(session)
    response.delete_cookie("session")
    return {"status": "ok"}


# ---------------------------------------------------------------------------
# Session proxy (forward to amplifierd)
# ---------------------------------------------------------------------------


@app.get("/api/sessions")
async def get_sessions(
    request: Request,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Proxy GET /sessions to amplifierd."""
    resp = await request.app.state.http_client.get("/sessions")
    return JSONResponse(content=resp.json(), status_code=resp.status_code)


@app.post("/api/sessions")
async def create_session(
    request: Request,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Proxy POST /sessions to amplifierd, injecting bundle_name if absent."""
    body: dict[str, Any] = await request.json()
    body.setdefault("bundle_name", BUNDLE_NAME)
    resp = await request.app.state.http_client.post("/sessions", json=body)
    return JSONResponse(content=resp.json(), status_code=resp.status_code)


@app.get("/api/sessions/{session_id}")
async def get_session(
    session_id: str,
    request: Request,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Proxy GET /sessions/{id} to amplifierd."""
    resp = await request.app.state.http_client.get(f"/sessions/{session_id}")
    return JSONResponse(content=resp.json(), status_code=resp.status_code)


@app.delete("/api/sessions/{session_id}")
async def delete_session(
    session_id: str,
    request: Request,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Proxy DELETE /sessions/{id} to amplifierd."""
    resp = await request.app.state.http_client.delete(f"/sessions/{session_id}")
    return JSONResponse(content=resp.json(), status_code=resp.status_code)


@app.get("/api/sessions/{session_id}/transcript")
async def get_transcript(
    session_id: str,
    request: Request,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Proxy GET /sessions/{id}/transcript to amplifierd."""
    resp = await request.app.state.http_client.get(f"/sessions/{session_id}/transcript")
    return JSONResponse(content=resp.json(), status_code=resp.status_code)


# ---------------------------------------------------------------------------
# AG-UI / CopilotKit streaming endpoint
# ---------------------------------------------------------------------------


@app.post("/api/copilotkit")
async def copilotkit(request: Request) -> StreamingResponse:
    """AG-UI streaming endpoint.

    Accepts CopilotKit protocol payload, executes via amplifierd, and streams
    AG-UI SSE events back to the client.
    """
    from ag_ui.core import RunStartedEvent

    body = await request.json()
    thread_id: str = body.get("threadId", str(uuid.uuid4()))
    run_id: str = body.get("runId", str(uuid.uuid4()))
    messages: list[dict[str, Any]] = body.get("messages", [])

    # Extract last user message as prompt
    prompt: str | None = None
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                prompt = content
            elif isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        prompt = block.get("text", "")
                        break
            break

    if not prompt:
        raise HTTPException(status_code=400, detail="Missing user message")

    session_id = thread_id
    client = request.app.state.http_client

    # Create session if it doesn't exist
    check_resp = await client.get(f"/sessions/{session_id}")
    if check_resp.status_code == 404:
        await client.post(
            "/sessions",
            json={"session_id": session_id, "bundle_name": BUNDLE_NAME},
        )

    # Start streaming execution
    exec_resp = await client.post(
        f"/sessions/{session_id}/execute/stream",
        json={"prompt": prompt},
    )
    if exec_resp.status_code != 202:
        raise HTTPException(status_code=502, detail="Failed to start execution")

    translator = AmplifierdEventTranslator(run_id=run_id, thread_id=thread_id)

    async def event_stream() -> AsyncGenerator[str, None]:
        # Yield RUN_STARTED first
        run_started = RunStartedEvent(thread_id=thread_id, run_id=run_id)
        yield _sse(run_started)

        sse_state: dict[str, Any] = {}
        async with client.stream("GET", f"/events?session={session_id}") as stream_resp:
            async for line in stream_resp.aiter_lines():
                parsed = _parse_sse(line, sse_state)
                if parsed:
                    sse_state = {}  # Reset state for next event
                    ag_events = translator.translate(parsed)
                    for ag_event in ag_events:
                        yield _sse(ag_event)
                    if parsed.get("event_type") == "orchestrator:complete":
                        break

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


# ---------------------------------------------------------------------------
# Artifact endpoints
# ---------------------------------------------------------------------------


@app.get("/api/artifacts/{session_id}")
async def list_session_artifacts(
    session_id: str,
    email: str = Depends(_require_auth),
) -> list[dict[str, Any]]:
    """Return sorted list of artifacts for a session."""
    return list_artifacts(session_id)


@app.get("/api/artifacts/{session_id}/{name}")
async def get_session_artifact(
    session_id: str,
    name: str,
    email: str = Depends(_require_auth),
) -> JSONResponse:
    """Return the content of a specific artifact."""
    content = get_artifact(session_id, name)
    if content is None:
        raise HTTPException(status_code=404, detail="Artifact not found")
    return JSONResponse({"content": content})


# ---------------------------------------------------------------------------
# MCP App serving
# ---------------------------------------------------------------------------


@app.get("/api/apps/{app_name}")
async def serve_mcp_app(
    app_name: str,
    email: str = Depends(_require_auth),
) -> HTMLResponse:
    """Serve an MCP App HTML file from the bundle/apps directory."""
    # Normalize: add .html extension if missing
    if not app_name.endswith(".html"):
        app_name = app_name + ".html"

    app_path = APPS_DIR / app_name
    if not app_path.exists():
        raise HTTPException(status_code=404, detail="App not found")

    return HTMLResponse(content=app_path.read_text(encoding="utf-8"))


# ---------------------------------------------------------------------------
# Static frontend — MUST be last (catch-all)
# ---------------------------------------------------------------------------

_FRONTEND_DIST = Path(__file__).parent.parent / "frontend" / "dist"
if _FRONTEND_DIST.exists():
    app.mount(
        "/",
        StaticFiles(directory=str(_FRONTEND_DIST), html=True),
        name="static",
    )