from __future__ import annotations

import time
import uuid
from contextlib import asynccontextmanager
from typing import Literal

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from .orchestrator import FuguLiteOrchestrator


class RouteRequest(BaseModel):
    prompt: str
    domain: str = "general"
    tags: list[str] = Field(default_factory=list)


class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant", "tool"]
    content: str


class ChatCompletionRequest(BaseModel):
    model: str = "fugu-lite"
    messages: list[ChatMessage]
    domain: str = "general"
    tags: list[str] = Field(default_factory=list)
    stream: bool = False


def create_app(checkpoint: str, worker_config: str) -> FastAPI:
    orchestrator = FuguLiteOrchestrator(checkpoint, worker_config)

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        yield
        await orchestrator.close()

    app = FastAPI(title="BenchGen Fugu-Lite", version="0.1.0", lifespan=lifespan)

    @app.get("/health")
    async def health():
        return {"status": "ok", "workers": orchestrator.pool.worker_ids}

    @app.post("/route")
    async def route(request: RouteRequest):
        return orchestrator.route(request.prompt, request.domain, request.tags)

    @app.post("/v1/chat/completions")
    async def chat_completions(request: ChatCompletionRequest):
        if request.stream:
            raise HTTPException(status_code=400, detail="Streaming is not implemented in v0.1")
        if not request.messages:
            raise HTTPException(status_code=400, detail="messages cannot be empty")
        flattened = "\n\n".join(
            f"{message.role.upper()}: {message.content}" for message in request.messages
        )
        try:
            answer = await orchestrator.answer(flattened, request.domain, request.tags)
        except Exception as exc:
            raise HTTPException(status_code=502, detail=str(exc)) from exc
        result = answer["result"]
        return {
            "id": f"chatcmpl-fugu-lite-{uuid.uuid4().hex}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": result["response"]},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": result.get("prompt_tokens"),
                "completion_tokens": result.get("completion_tokens"),
                "total_tokens": (
                    (result.get("prompt_tokens") or 0) + (result.get("completion_tokens") or 0)
                ),
            },
            "fugu_lite": {
                "routing": answer["routing"],
                "served_model": result.get("served_model"),
                "cost_usd": result.get("cost_usd"),
                "latency_ms": result.get("latency_ms"),
                "fallback_errors": answer["fallback_errors"],
            },
        }

    return app