from __future__ import annotations

import json
import os
import secrets
import threading
import time
import uuid
from contextlib import asynccontextmanager
from typing import Any, Literal

from fastapi import Depends, FastAPI, Header, HTTPException
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from tokenizers import Tokenizer

from nayara.artifacts import materialize_checkpoint
from nayara.generate import load_model, stream_generate


class CompletionRequest(BaseModel):
    model: str = "nayara-v1"
    prompt: str
    max_tokens: int = Field(192, ge=1, le=512)
    temperature: float = Field(0.2, ge=0, le=2)
    top_k: int = Field(40, ge=1, le=32000)
    stream: bool = False


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


class ChatRequest(BaseModel):
    model: str = "nayara-v1"
    messages: list[ChatMessage]
    max_tokens: int = Field(192, ge=1, le=512)
    temperature: float = Field(0.2, ge=0, le=2)
    top_k: int = Field(40, ge=1, le=32000)
    stream: bool = False


runtime: dict[str, Any] = {"ready": False}
generation_lock = threading.Lock()
MODEL_ID = os.environ.get("NAYARA_MODEL_ID", "nayara-v1")


def _authorize(authorization: str | None = Header(default=None)) -> None:
    expected = os.environ.get("NAYARA_API_KEY")
    if not expected:
        return
    supplied = authorization.removeprefix("Bearer ") if authorization else ""
    if not secrets.compare_digest(supplied, expected):
        raise HTTPException(status_code=401, detail="invalid bearer token")


def _chat_prompt(messages: list[ChatMessage]) -> str:
    if not messages or not any(message.role == "user" for message in messages):
        raise HTTPException(status_code=422, detail="messages must contain at least one user message")
    parts: list[str] = []
    for message in messages:
        if message.role == "system":
            parts.append(f"<|user|>\nSystem guidance: {message.content.strip()}")
        else:
            parts.append(f"<|{message.role}|>\n{message.content.strip()}")
    parts.append("<|assistant|>\n")
    return "\n".join(parts)


def _generate(prompt: str, request: CompletionRequest | ChatRequest):
    return stream_generate(
        runtime["model"], runtime["tokenizer"], prompt,
        max_new_tokens=request.max_tokens, temperature=request.temperature, top_k=request.top_k,
    )


def _complete(prompt: str, request: CompletionRequest | ChatRequest) -> str:
    final = ""
    with generation_lock:
        for final in _generate(prompt, request):
            pass
    return final


def _sse(prompt: str, request: CompletionRequest | ChatRequest, chat: bool):
    request_id = f"nayara-{uuid.uuid4().hex}"
    previous = ""
    with generation_lock:
        for partial in _generate(prompt, request):
            delta, previous = partial[len(previous) :], partial
            if chat:
                payload = {"id": request_id, "object": "chat.completion.chunk", "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}]}
            else:
                payload = {"id": request_id, "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": delta, "finish_reason": None}]}
            yield f"data: {json.dumps(payload)}\n\n"
    yield "data: [DONE]\n\n"


@asynccontextmanager
async def lifespan(_: FastAPI):
    if os.environ.get("NAYARA_REQUIRE_API_KEY") == "1" and not os.environ.get("NAYARA_API_KEY"):
        raise RuntimeError("NAYARA_REQUIRE_API_KEY=1 but NAYARA_API_KEY is unset")
    checkpoint = await run_in_threadpool(materialize_checkpoint)
    threads = int(os.environ.get("NAYARA_THREADS", max(1, (os.cpu_count() or 4) // 2)))
    runtime["tokenizer"] = Tokenizer.from_file(str(checkpoint / "tokenizer.json"))
    runtime["model"] = await run_in_threadpool(load_model, checkpoint, True, threads)
    runtime.update({"ready": True, "checkpoint": str(checkpoint), "loaded_at": time.time()})
    yield
    runtime.clear()
    runtime["ready"] = False


app = FastAPI(title="NAYARA API", version="1.0.0", lifespan=lifespan)


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


@app.get("/readyz")
def ready() -> dict[str, Any]:
    if not runtime.get("ready"):
        raise HTTPException(status_code=503, detail="model is not loaded")
    return {"status": "ready", "model": MODEL_ID, "loaded_at": runtime["loaded_at"]}


@app.get("/v1/models", dependencies=[Depends(_authorize)])
def models() -> dict[str, Any]:
    return {"object": "list", "data": [{"id": MODEL_ID, "object": "model", "owned_by": "nayara"}]}


@app.post("/v1/completions", dependencies=[Depends(_authorize)])
async def completions(request: CompletionRequest):
    prompt = request.prompt
    if request.stream:
        return StreamingResponse(_sse(prompt, request, False), media_type="text/event-stream")
    text = await run_in_threadpool(_complete, prompt, request)
    return {"id": f"nayara-{uuid.uuid4().hex}", "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": text, "finish_reason": "stop"}]}


@app.post("/v1/chat/completions", dependencies=[Depends(_authorize)])
async def chat_completions(request: ChatRequest):
    prompt = _chat_prompt(request.messages)
    if request.stream:
        return StreamingResponse(_sse(prompt, request, True), media_type="text/event-stream")
    text = await run_in_threadpool(_complete, prompt, request)
    return {"id": f"nayara-{uuid.uuid4().hex}", "object": "chat.completion", "model": MODEL_ID, "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}]}