"""LLM Gateway proxy routes — forwards client requests to LLM providers.

Supports both non-streaming (`POST /api/v1/llm/chat`) and SSE streaming
(`POST /api/v1/llm/chat/stream`) modes.
"""

import json
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, Field

from agentkit.core.exceptions import LLMProviderError, ModelNotFoundError
from agentkit.llm.gateway import QuotaExceededError
from agentkit.llm.protocol import LLMResponse, StreamChunk, TokenUsage, ToolCall
from agentkit.llm.providers.usage_store import UsageStoreUnavailableError
from agentkit.server.admin.context import get_department_context

router = APIRouter(prefix="/llm", tags=["llm-gateway"])


class LLMChatRequest(BaseModel):
    """Request body for LLM chat endpoints."""

    model_config = ConfigDict(extra="forbid")

    messages: list[dict[str, str]] = Field(..., min_length=1)
    model: str
    temperature: float = 0.7
    max_tokens: int = 2000
    tools: list[dict[str, Any]] | None = None
    tool_choice: str = "auto"
    timeout: float | None = None


def _serialize_tool_call(tc: ToolCall) -> dict[str, Any]:
    return {"id": tc.id, "name": tc.name, "arguments": tc.arguments}


def _serialize_usage(usage: TokenUsage) -> dict[str, int]:
    return {
        "prompt_tokens": usage.prompt_tokens,
        "completion_tokens": usage.completion_tokens,
        "total_tokens": usage.total_tokens,
    }


def _serialize_response(response: LLMResponse) -> dict[str, Any]:
    return {
        "content": response.content,
        "model": response.model,
        "usage": _serialize_usage(response.usage),
        "tool_calls": [_serialize_tool_call(tc) for tc in response.tool_calls],
        "latency_ms": response.latency_ms,
    }


def _serialize_chunk(chunk: StreamChunk) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "content": chunk.content,
        "model": chunk.model,
        "is_final": chunk.is_final,
    }
    if chunk.tool_calls:
        payload["tool_calls"] = [_serialize_tool_call(tc) for tc in chunk.tool_calls]
    if chunk.usage is not None:
        payload["usage"] = _serialize_usage(chunk.usage)
    return payload


def _quota_error_payload(exc: QuotaExceededError) -> dict[str, Any]:
    """Build a structured 429 error body from a QuotaExceededError."""
    return {
        "error": "quota_exceeded",
        "department_id": exc.department_id,
        "quota_type": exc.quota_type,
        "period": exc.period,
        "limit": exc.limit,
        "current": exc.current,
    }


@router.post("/chat")
async def chat(
    request: Request,
    body: LLMChatRequest,
    ctx: Any = Depends(get_department_context),
) -> dict[str, Any]:
    """Non-streaming LLM chat proxy.

    Forwards the request to the configured LLMGateway and returns the
    serialized LLMResponse. Quota-exceeded errors from the gateway are
    translated to HTTP 429.
    """
    gateway = request.app.state.llm_gateway
    db_path = getattr(request.app.state, "auth_db_path", None)
    try:
        response = await gateway.chat(
            messages=body.messages,
            model=body.model,
            tools=body.tools,
            tool_choice=body.tool_choice,
            timeout=body.timeout,
            temperature=body.temperature,
            max_tokens=body.max_tokens,
            user_id=ctx.user_id,
            department_ids=ctx.department_ids if ctx.department_ids else None,
            db_path=db_path,
        )
    except QuotaExceededError as e:
        raise HTTPException(status_code=429, detail=_quota_error_payload(e)) from e
    except UsageStoreUnavailableError as e:
        raise HTTPException(
            status_code=503,
            detail={"error": "usage_store_unavailable", "detail": str(e)},
        ) from e
    except ModelNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e)) from e
    except LLMProviderError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    return _serialize_response(response)


@router.post("/chat/stream")
async def chat_stream(
    request: Request,
    body: LLMChatRequest,
    ctx: Any = Depends(get_department_context),
) -> StreamingResponse:
    """SSE streaming LLM chat proxy.

    Each StreamChunk is serialized as `data: {json}\\n\\n`. The stream
    terminates with `data: [DONE]\\n\\n`.
    """

    async def event_generator():
        gateway = request.app.state.llm_gateway
        db_path = getattr(request.app.state, "auth_db_path", None)
        try:
            async for chunk in gateway.chat_stream(
                messages=body.messages,
                model=body.model,
                tools=body.tools,
                tool_choice=body.tool_choice,
                timeout=body.timeout,
                temperature=body.temperature,
                max_tokens=body.max_tokens,
                user_id=ctx.user_id,
                department_ids=ctx.department_ids if ctx.department_ids else None,
                db_path=db_path,
            ):
                payload = _serialize_chunk(chunk)
                yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
        except QuotaExceededError as e:
            error_payload = _quota_error_payload(e)
            error_payload["error"] = "quota_exceeded"
            yield f"data: {json.dumps(error_payload, ensure_ascii=False)}\n\n"
        except UsageStoreUnavailableError as e:
            error_payload = {
                "error": "usage_store_unavailable",
                "detail": str(e),
            }
            yield f"data: {json.dumps(error_payload, ensure_ascii=False)}\n\n"
        except ModelNotFoundError as e:
            error_payload = {"error": "model_not_found", "detail": str(e)}
            yield f"data: {json.dumps(error_payload, ensure_ascii=False)}\n\n"
        except LLMProviderError as e:
            error_payload = {"error": "provider_error", "detail": str(e)}
            yield f"data: {json.dumps(error_payload, ensure_ascii=False)}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        content=event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )
