"""OpenAI-compatible THOX interactive and specialist model service.

The interactive model is intentionally small and eagerly loaded so readiness
means a user request can start generating immediately.  The 25B Rust-specialist
model remains lazy because loading it takes roughly 100 seconds on cpu-upgrade
and it is not part of the interactive latency contract.
"""

from __future__ import annotations

import json
import os
import queue
import threading
import time
import uuid
from contextlib import asynccontextmanager
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from huggingface_hub import hf_hub_download
from pydantic import BaseModel, ConfigDict, Field, field_validator

FAST_MODEL_ID = "thox-fast-chat"
FAST_MODEL_REPO = os.environ.get(
    "THOX_FAST_MODEL_REPO", "Qwen/Qwen2.5-0.5B-Instruct-GGUF"
)
FAST_MODEL_REVISION = os.environ.get(
    "THOX_FAST_MODEL_REVISION", "9217f5db79a29953eb74d5343926648285ec7e67"
)
FAST_MODEL_FILE = os.environ.get(
    "THOX_FAST_MODEL_FILE", "qwen2.5-0.5b-instruct-q4_k_m.gguf"
)
CODER_MODEL_ID = "thox-rust-coder"
CODER_MODEL_REPO = os.environ.get(
    "THOX_CODER_MODEL_REPO", "Em-80/Qwen3-coder-REAP-25B-A3B-Rust-GGUF"
)
CODER_MODEL_REVISION = os.environ.get("THOX_CODER_MODEL_REVISION", "main")
CODER_MODEL_FILE = os.environ.get(
    "THOX_CODER_MODEL_FILE", "Qwen3-Coder-REAP-25B-A3B-Rust-Q4_K_M.gguf"
)
N_CTX = int(os.environ.get("THOX_N_CTX", "4096"))
FAST_POOL_SIZE = max(1, min(int(os.environ.get("THOX_FAST_POOL_SIZE", "1")), 2))
FAST_QUEUE_TIMEOUT_S = max(
    0.1, min(float(os.environ.get("THOX_FAST_QUEUE_TIMEOUT_S", "6")), 10.0)
)
MAX_OUTPUT_TOKENS = max(
    1, min(int(os.environ.get("THOX_MAX_OUTPUT_TOKENS", "128")), 512)
)
FAST_MAX_OUTPUT_TOKENS = max(
    1, min(int(os.environ.get("THOX_FAST_MAX_OUTPUT_TOKENS", "16")), 32)
)
STREAM_HEARTBEAT_S = max(
    0.25, min(float(os.environ.get("THOX_STREAM_HEARTBEAT_S", "2")), 5.0)
)
MAX_MESSAGES = 64
MAX_MESSAGE_CHARS = 65_536
MAX_REQUEST_CHARS = 131_072

PROVIDER_SYSTEM_PREFIX = (
    "Follow the caller-provided system instructions and requested output format "
    "exactly. Treat later system messages as the authoritative assistant persona "
    "and identity. Never claim that cloud inference ran locally or on-device."
)
DEFAULT_SYSTEM = (
    "You are THOX Fast Chat, a concise privacy-first assistant. Follow the "
    "user's requested output format exactly. Never claim that cloud inference "
    "ran locally or on-device."
)

SOURCE_REVISION_FILE = Path(
    os.environ.get("THOX_SOURCE_REVISION_FILE", "/app/THOXROUTE_GITHUB_SHA")
)


def _source_revision() -> str:
    """Return the exact owning Git revision baked into the Space image."""

    try:
        revision = SOURCE_REVISION_FILE.read_text(encoding="utf-8").strip()
    except OSError:
        return "unknown"
    return revision if len(revision) == 40 and all(c in "0123456789abcdef" for c in revision) else "unknown"


def _usable_cpus() -> int:
    """Return the container CPU quota instead of the misleading host count."""

    try:
        quota, period = open("/sys/fs/cgroup/cpu.max", encoding="utf-8").read().split()
        if quota != "max":
            return max(1, int(int(quota) / int(period)))
    except (OSError, ValueError):
        pass
    try:
        quota = int(
            open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us", encoding="utf-8").read()
        )
        period = int(
            open("/sys/fs/cgroup/cpu/cpu.cfs_period_us", encoding="utf-8").read()
        )
        if quota > 0:
            return max(1, quota // period)
    except (OSError, ValueError):
        pass
    try:
        return max(1, len(os.sched_getaffinity(0)))
    except (AttributeError, OSError):
        return os.cpu_count() or 2


def _download(repo: str, filename: str, revision: str) -> str:
    return hf_hub_download(
        repo_id=repo,
        filename=filename,
        revision=revision,
        token=os.environ.get("HF_TOKEN") or None,
    )


def _new_llama(path: str, *, threads: int):
    from llama_cpp import Llama

    return Llama(
        model_path=path,
        n_ctx=N_CTX,
        n_threads=threads,
        # llama.cpp uses the batch pool for prompt evaluation and may derive
        # its default from the host CPU count. Spaces expose more host CPUs
        # than the container quota, so leaving this unset oversubscribes the
        # exact cold-prefix phase that owns the interactive latency budget.
        n_threads_batch=threads,
        n_gpu_layers=int(os.environ.get("THOX_GPU_LAYERS", "-1")),
        verbose=False,
    )


def _warm_interactive_model(model: Any) -> None:
    """Pay one-time graph/template initialization before readiness is visible."""

    model.create_chat_completion(
        messages=[
            {"role": "system", "content": PROVIDER_SYSTEM_PREFIX},
            {"role": "system", "content": DEFAULT_SYSTEM},
            {"role": "user", "content": "Reply with OK."},
        ],
        max_tokens=1,
        temperature=0.0,
    )


@dataclass
class Lease:
    model: Any
    release: Any


class Runtime:
    """Own bounded model capacity without sharing one llama context concurrently."""

    def __init__(self) -> None:
        fast_path = _download(FAST_MODEL_REPO, FAST_MODEL_FILE, FAST_MODEL_REVISION)
        usable = _usable_cpus()
        per_model_threads = max(1, usable // FAST_POOL_SIZE)
        self._fast: queue.LifoQueue[Any] = queue.LifoQueue(maxsize=FAST_POOL_SIZE)
        for _ in range(FAST_POOL_SIZE):
            model = _new_llama(fast_path, threads=per_model_threads)
            _warm_interactive_model(model)
            self._fast.put(model)
        self._coder = None
        self._coder_lock = threading.Lock()

    def acquire(self, model_id: str) -> Lease:
        if model_id == FAST_MODEL_ID:
            try:
                model = self._fast.get(timeout=FAST_QUEUE_TIMEOUT_S)
            except queue.Empty as exc:
                raise HTTPException(status_code=429, detail="interactive capacity busy") from exc
            return Lease(model=model, release=lambda: self._fast.put(model))
        if model_id == CODER_MODEL_ID:
            if not self._coder_lock.acquire(blocking=False):
                raise HTTPException(status_code=429, detail="specialist capacity busy")
            try:
                if self._coder is None:
                    path = _download(
                        CODER_MODEL_REPO, CODER_MODEL_FILE, CODER_MODEL_REVISION
                    )
                    self._coder = _new_llama(path, threads=_usable_cpus())
            except Exception:
                self._coder_lock.release()
                raise
            return Lease(model=self._coder, release=self._coder_lock.release)
        raise HTTPException(status_code=404, detail="model not found")

    @property
    def interactive_available(self) -> int:
        return self._fast.qsize()


class Msg(BaseModel):
    model_config = ConfigDict(extra="forbid")

    role: str = Field(pattern=r"^(system|user|assistant|tool)$")
    content: str = Field(min_length=1, max_length=MAX_MESSAGE_CHARS)


class ChatRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    model: str = FAST_MODEL_ID
    messages: list[Msg] = Field(min_length=1, max_length=MAX_MESSAGES)
    max_tokens: int = Field(default=64, ge=1, le=4096)
    temperature: float = Field(default=0.2, ge=0.0, le=2.0)
    stream: bool = False

    @field_validator("messages")
    @classmethod
    def bound_aggregate_prompt(cls, messages: list[Msg]) -> list[Msg]:
        if sum(len(message.content) for message in messages) > MAX_REQUEST_CHARS:
            raise ValueError("aggregate prompt is too large")
        return messages


runtime: Runtime | None = None


@asynccontextmanager
async def lifespan(_: FastAPI):
    global runtime
    runtime = Runtime()
    try:
        yield
    finally:
        runtime = None


api = FastAPI(
    title="THOX interactive model service", version="1.0.0", lifespan=lifespan
)


def _messages(req: ChatRequest) -> list[dict[str, str]]:
    """Build a bounded prompt whose first tokens survive cross-persona traffic.

    Every fast-chat request begins with the exact prefix paid for during startup
    warmup. Caller system messages retain their order immediately after it and
    remain authoritative, while callers without a system message receive the
    existing default assistant persona. Specialist prompts keep their previous
    behavior and never inherit the fast-chat prefix.
    """

    messages = [message.model_dump() for message in req.messages]
    has_caller_system = any(message["role"] == "system" for message in messages)
    if req.model == FAST_MODEL_ID:
        provider_prefix = {"role": "system", "content": PROVIDER_SYSTEM_PREFIX}
        if not messages or messages[0] != provider_prefix:
            messages.insert(0, provider_prefix)
        if not has_caller_system:
            messages.insert(1, {"role": "system", "content": DEFAULT_SYSTEM})
    elif not has_caller_system:
        messages.insert(0, {"role": "system", "content": DEFAULT_SYSTEM})
    return messages


def _max_tokens(req: ChatRequest) -> int:
    ceiling = FAST_MAX_OUTPUT_TOKENS if req.model == FAST_MODEL_ID else MAX_OUTPUT_TOKENS
    return min(req.max_tokens, ceiling)


def _runtime() -> Runtime:
    if runtime is None:
        raise HTTPException(status_code=503, detail="model runtime is not ready")
    return runtime


@api.get("/")
@api.get("/healthz")
def healthz() -> dict[str, Any]:
    active = _runtime()
    return {
        "status": "ready",
        "interactive_model": FAST_MODEL_ID,
        "interactive_model_repo": FAST_MODEL_REPO,
        "interactive_model_revision": FAST_MODEL_REVISION,
        "interactive_pool_size": FAST_POOL_SIZE,
        "interactive_available": active.interactive_available,
        "interactive_common_prefix_warmed": True,
        "specialist_model": CODER_MODEL_ID,
        "interactive_max_output_tokens": FAST_MAX_OUTPUT_TOKENS,
        "stream_heartbeat_s": STREAM_HEARTBEAT_S,
        "max_output_tokens": MAX_OUTPUT_TOKENS,
        "n_ctx": N_CTX,
        "threads": _usable_cpus(),
        "prompt_threads": _usable_cpus(),
        "source_revision": _source_revision(),
    }


def _completion_id() -> str:
    return "chatcmpl-" + uuid.uuid4().hex[:16]


def _stream_completion(req: ChatRequest, lease: Lease) -> Iterator[bytes]:
    completion_id = _completion_id()
    created = int(time.time())
    items: queue.Queue[dict[str, Any] | object] = queue.Queue()
    complete = object()
    failed = object()

    def generate() -> None:
        try:
            chunks = lease.model.create_chat_completion(
                messages=_messages(req),
                max_tokens=_max_tokens(req),
                temperature=req.temperature,
                stream=True,
            )
            for chunk in chunks:
                items.put(chunk)
            items.put(complete)
        except Exception:
            # Provider-controlled exception text must never cross the API boundary.
            items.put(failed)
        finally:
            # A disconnected client closes the response generator, but llama.cpp
            # can still be using the context. The worker therefore owns release.
            lease.release()

    threading.Thread(target=generate, daemon=True, name="thox-fast-generation").start()

    # Prove liveness before prompt evaluation, then emit SSE comments while the
    # synchronous llama context is busy. Upstream read timers reset on every
    # heartbeat and comments are ignored by OpenAI-compatible parsers.
    initial = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": req.model,
        "choices": [
            {
                "index": 0,
                "delta": {"role": "assistant", "content": ""},
                "finish_reason": None,
            }
        ],
    }
    yield f"data: {json.dumps(initial, separators=(',', ':'))}\n\n".encode()
    while True:
        try:
            item = items.get(timeout=STREAM_HEARTBEAT_S)
        except queue.Empty:
            yield b": thox-fast-heartbeat\n\n"
            continue
        if item is complete:
            yield b"data: [DONE]\n\n"
            return
        if item is failed:
            raise RuntimeError("interactive generation failed")
        if isinstance(item, dict):
            chunk = item
            choice = chunk.get("choices", [{}])[0]
            payload = {
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": req.model,
                "choices": [
                    {
                        "index": 0,
                        "delta": choice.get("delta") or {},
                        "finish_reason": choice.get("finish_reason"),
                    }
                ],
            }
            yield f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()


@api.post("/v1/chat/completions")
def chat_completions(req: ChatRequest):
    lease = _runtime().acquire(req.model)
    if req.stream:
        return StreamingResponse(
            _stream_completion(req, lease),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    started = time.monotonic()
    try:
        result = lease.model.create_chat_completion(
            messages=_messages(req),
            max_tokens=_max_tokens(req),
            temperature=req.temperature,
        )
    finally:
        lease.release()
    elapsed = time.monotonic() - started
    choice = result["choices"][0]
    usage = result.get("usage") or {}
    return {
        "id": _completion_id(),
        "object": "chat.completion",
        "created": int(time.time()),
        "model": req.model,
        "choices": [
            {
                "index": 0,
                "finish_reason": choice.get("finish_reason") or "stop",
                "message": choice["message"],
            }
        ],
        "usage": {
            "prompt_tokens": int(usage.get("prompt_tokens") or 0),
            "completion_tokens": int(usage.get("completion_tokens") or 0),
            "total_tokens": int(usage.get("total_tokens") or 0),
        },
        "thox_perf": {"seconds": round(elapsed, 3)},
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(api, host="0.0.0.0", port=int(os.environ.get("PORT", "7860")))