"""GGUF LLM API for Hugging Face Spaces (Docker SDK) — runtime-swappable models.

Downloads a .gguf from a HF model repo, loads it with llama-cpp-python and
serves an OpenAI-compatible API. The initial model comes from environment
variables; afterwards a NEW model can be loaded at runtime via POST /admin/load
— no Space variable change, therefore NO image rebuild.

Environment variables (initial model only):
    REPO_ID, FILENAME, HF_TOKEN, MODEL_ID, N_CTX, N_THREADS, N_GPU_LAYERS,
    CHAT_FORMAT ("" = auto from GGUF metadata), DEFAULT_MAX_TOKENS
"""
import os
import time
import uuid
import logging
import threading
from contextlib import asynccontextmanager
from typing import List, Optional, Dict, Any

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from llama_cpp.llama_cache import LlamaRAMCache
from llama_cpp import llama_chat_format

logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)s  %(message)s")
log = logging.getLogger("gguf-api")


# DeepSeek-LLM chat template — not built into llama-cpp-python, and TheBloke's
# 2023 GGUFs predate embedded chat_template metadata. Registered here so this
# Space can keep serving deepseek-7b-chat (its default model).
@llama_chat_format.register_chat_format("deepseek")
def _format_deepseek(messages, **kwargs) -> llama_chat_format.ChatFormatterResponse:
    prompt = ""
    for m in messages:
        role, content = m["role"], (m["content"] or "").strip()
        if role == "system":
            prompt += content + "\n\n"
        elif role == "user":
            prompt += f"User: {content}\n\n"
        elif role == "assistant":
            prompt += f"Assistant: {content}<｜end▁of▁sentence｜>"
    prompt += "Assistant:"
    return llama_chat_format.ChatFormatterResponse(prompt=prompt, stop=["User:"])

HF_TOKEN           = os.getenv("HF_TOKEN") or None
N_CTX              = int(os.getenv("N_CTX", "4096"))
_USABLE = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 4)
N_THREADS          = int(os.getenv("N_THREADS", str(_USABLE)))
N_GPU_LAYERS       = int(os.getenv("N_GPU_LAYERS", "0"))
DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "512"))
# Prompt (KV) cache. Without it every call re-prefills the whole system prompt,
# which is NOT cheap on 2 vCPU: measured 20.8s -> 8.5s per call with the study
# plan's rotating daily/weekly/monthly prompts, ~21% off a full plan.
PROMPT_CACHE_MB    = int(os.getenv("PROMPT_CACHE_MB", "512"))   # 0 disables

# Current model state (mutable at runtime via /admin/load).
_llm: Optional[Llama] = None
_llm_lock = threading.Lock()
_load_error: Optional[str] = None
_loading: bool = False
_cfg: Dict[str, str] = {
    "repo_id":     os.getenv("REPO_ID", "TheBloke/deepseek-llm-7B-chat-GGUF").strip(),
    "filename":    os.getenv("FILENAME", "deepseek-llm-7b-chat.Q4_K_M.gguf").strip(),
    "model_id":    (os.getenv("MODEL_ID", "") or "deepseek-7b-chat").strip(),
    "chat_format": os.getenv("CHAT_FORMAT", "deepseek").strip(),
}


def _download_and_load(cfg: Dict[str, str]) -> None:
    """Download cfg's .gguf and load it, replacing the current model."""
    global _llm, _load_error, _loading, _cfg
    _loading = True
    _load_error = None
    try:
        log.info("Downloading %s from %s ...", cfg["filename"], cfg["repo_id"])
        t0 = time.time()
        path = hf_hub_download(repo_id=cfg["repo_id"], filename=cfg["filename"],
                               token=HF_TOKEN)
        log.info("Downloaded in %.1fs; loading ...", time.time() - t0)
        kwargs: Dict[str, Any] = dict(model_path=path, n_ctx=N_CTX,
                                      n_threads=N_THREADS, n_gpu_layers=N_GPU_LAYERS,
                                      verbose=False)
        if cfg["chat_format"]:
            kwargs["chat_format"] = cfg["chat_format"]
        new_llm = Llama(**kwargs)
        if PROMPT_CACHE_MB > 0:
            new_llm.set_cache(LlamaRAMCache(capacity_bytes=PROMPT_CACHE_MB * 1024 * 1024))
            log.info("Prompt cache enabled (%d MB).", PROMPT_CACHE_MB)
        with _llm_lock:                      # swap atomically; old model is freed
            _llm = new_llm
            _cfg = dict(cfg)
        log.info("Model '%s' ready.", cfg["model_id"])
    except Exception as exc:
        _load_error = f"Failed to load {cfg['model_id']}: {exc}"
        log.exception(_load_error)
    finally:
        _loading = False


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the initial model in the background so port 7860 opens immediately.
    threading.Thread(target=_download_and_load, args=(dict(_cfg),),
                     name="model-loader", daemon=True).start()
    yield


app = FastAPI(title="GGUF LLM API (runtime-swappable)", version="2.0.0", lifespan=lifespan)


class ChatMessage(BaseModel):
    role: str
    content: str = ""


class ChatCompletionRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    temperature: float = 0.7
    top_p: float = 0.95
    max_tokens: Optional[int] = None
    stop: Optional[List[str]] = None
    response_format: Optional[Dict[str, Any]] = None


class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = Field(default=512)
    temperature: float = 0.7


class LoadRequest(BaseModel):
    repo_id: str
    filename: str
    model_id: str
    chat_format: str = ""     # "" = auto from GGUF metadata


def _require_model() -> Llama:
    if _llm is None or _loading:
        raise HTTPException(503, _load_error or "Model is loading, try again shortly.")
    return _llm


@app.get("/health")
def health() -> Dict[str, Any]:
    return {
        "status": "ok" if (_llm is not None and not _loading) else "loading_or_error",
        "model": _cfg["model_id"], "repo_id": _cfg["repo_id"],
        "filename": _cfg["filename"], "chat_format": _cfg["chat_format"],
        "loading": _loading, "n_ctx": N_CTX, "n_threads": N_THREADS,
        "n_gpu_layers": N_GPU_LAYERS, "error": _load_error,
    }


@app.post("/admin/load")
def admin_load(req: LoadRequest) -> Dict[str, Any]:
    """Load a different GGUF at runtime (async; poll /health until status ok)."""
    global _loading
    if _loading:
        raise HTTPException(409, "another model is already loading")
    cfg = {"repo_id": req.repo_id.strip(), "filename": req.filename.strip(),
           "model_id": req.model_id.strip(), "chat_format": req.chat_format.strip()}
    threading.Thread(target=_download_and_load, args=(cfg,),
                     name="model-loader", daemon=True).start()
    return {"status": "loading", "model": cfg["model_id"]}


@app.get("/v1/models")
def list_models() -> Dict[str, Any]:
    return {"object": "list",
            "data": [{"id": _cfg["model_id"], "object": "model", "owned_by": "local"}]}


@app.post("/v1/chat/completions")
def chat_completions(req: ChatCompletionRequest) -> Dict[str, Any]:
    llm = _require_model()
    if not req.messages:
        raise HTTPException(400, "`messages` must not be empty.")
    kwargs: Dict[str, Any] = dict(
        messages=[m.model_dump() for m in req.messages],
        temperature=req.temperature, top_p=req.top_p,
        max_tokens=req.max_tokens if req.max_tokens is not None else DEFAULT_MAX_TOKENS,
        stop=req.stop,
    )
    if req.response_format and req.response_format.get("type") == "json_object":
        kwargs["response_format"] = {"type": "json_object"}
    try:
        with _llm_lock:
            result = llm.create_chat_completion(**kwargs)
    except Exception as exc:
        raise HTTPException(500, f"generation failed: {exc}") from exc
    result["id"] = result.get("id") or f"chatcmpl-{uuid.uuid4().hex}"
    result["model"] = _cfg["model_id"]
    result["created"] = result.get("created") or int(time.time())
    return result


@app.post("/generate")
def generate(req: GenerateRequest) -> Dict[str, Any]:
    llm = _require_model()
    t0 = time.time()
    try:
        with _llm_lock:
            result = llm.create_chat_completion(
                messages=[{"role": "user", "content": req.prompt}],
                temperature=req.temperature, max_tokens=req.max_tokens)
    except Exception as exc:
        raise HTTPException(500, f"generation failed: {exc}") from exc
    usage = result.get("usage", {}) or {}
    return {"prompt": req.prompt,
            "response": result["choices"][0]["message"]["content"],
            "model": _cfg["model_id"],
            "time_seconds": round(time.time() - t0, 2),
            "output_tokens": usage.get("completion_tokens")}