"""Local llama.cpp backend (GGUF). Self-contained, no network at inference time.

This is the default backend so the public HF Space stays self-hosted on CPU
(badges "Off the Grid" + "Llama Champion"). It serves TWO roles for TinySOC:

  * complete(messages)            -> the small explainer model (<=4B), chat JSON.
  * prompt_token_nll(ctx, line)   -> per-token negative log-likelihood of a log
                                     line, used by the perplexity highlighter.

A public Space cannot reach an external Ollama server, so that backend is
dev-only; the perplexity highlighter therefore needs prompt logprobs
locally. llama.cpp gives them in a SINGLE pass via `echo=True` + `logprobs`
(loading with logits_all=True), which is simpler than the Ollama forced-decode
hack. Lazy imports keep an Ollama-only dev box from needing llama-cpp-python.
"""
import os
from functools import lru_cache
from pathlib import Path

# --- Explainer model (chat) -------------------------------------------------
MODEL_PATH = os.environ.get("WEC_MODEL_PATH", "")
MODEL_REPO = os.environ.get("WEC_MODEL_REPO", "Mroqui/TinySOC-Qwen2.5-3B")
MODEL_FILE = os.environ.get("WEC_MODEL_FILE", "TinySOC-Qwen2.5-3B-Q5_K_M.gguf")
N_CTX = int(os.environ.get("WEC_N_CTX", "4096"))
MAX_TOKENS = int(os.environ.get("WEC_MAX_TOKENS", "768"))
N_GPU_LAYERS = int(os.environ.get("WEC_N_GPU_LAYERS", "0"))
TEMPERATURE = float(os.environ.get("WEC_TEMPERATURE", "0.2"))

# Preferred local explainer: the fine-tuned v3 GGUF ("Well-Tuned" badge).
_EXPLAIN_LOCAL = Path(__file__).parent / "models" / "TinySOC-Qwen2.5-3B-Q5_K_M.gguf"

# --- Perplexity model (highlighter) -----------------------------------------
# A small BASE model highlights novel tokens best. If none is provided we reuse
# the explainer model: detection still rests on the deterministic baseline, so
# the highlighter degrading to an instruct model never breaks correctness.
PPL_MODEL_PATH = os.environ.get("WEC_PPL_MODEL_PATH", "")
PPL_REPO = os.environ.get("WEC_PPL_REPO", "mradermacher/Qwen2.5-Coder-1.5B-GGUF")
PPL_FILE = os.environ.get("WEC_PPL_FILE", "Qwen2.5-Coder-1.5B.Q4_K_M.gguf")
PPL_N_CTX = int(os.environ.get("WEC_PPL_N_CTX", "1024"))
FLOOR_NLL = float(os.environ.get("WEC_FLOOR_NLL", "12.0"))


@lru_cache(maxsize=1)
def _explain_model_path() -> str:
    """Resolve the explainer GGUF: explicit env > local v3 > HF download."""
    if MODEL_PATH:
        return MODEL_PATH
    if _EXPLAIN_LOCAL.is_file():
        return str(_EXPLAIN_LOCAL)
    from huggingface_hub import hf_hub_download

    return hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)


@lru_cache(maxsize=1)
def _ppl_model_path() -> str:
    """Resolve the perplexity GGUF: explicit > HF download > reuse explainer."""
    if PPL_MODEL_PATH:
        return PPL_MODEL_PATH
    if PPL_REPO and PPL_FILE:
        from huggingface_hub import hf_hub_download

        return hf_hub_download(repo_id=PPL_REPO, filename=PPL_FILE)
    return _explain_model_path()


@lru_cache(maxsize=2)
def _load(model_path: str, n_ctx: int, logits_all: bool):
    """Load a GGUF once. Same file path => llama.cpp mmap shares weight pages."""
    from llama_cpp import Llama

    return Llama(
        model_path=model_path,
        n_ctx=n_ctx,
        n_threads=os.cpu_count() or 4,
        n_gpu_layers=N_GPU_LAYERS,
        logits_all=logits_all,
        verbose=False,
    )


def get_model():
    """Chat/explainer model instance (logits for the last token only)."""
    return _load(_explain_model_path(), N_CTX, False)


def get_ppl_model():
    """Perplexity model instance (logits_all=True so every prompt token scores)."""
    return _load(_ppl_model_path(), PPL_N_CTX, True)


def complete(messages: list[dict[str, str]]) -> str:
    """Generate the triage JSON via llama.cpp and return the raw text."""
    response = get_model().create_chat_completion(
        messages=messages,
        temperature=TEMPERATURE,
        max_tokens=MAX_TOKENS,
        response_format={"type": "json_object"},
    )
    return response["choices"][0]["message"]["content"]


def _nll(logits, token_id: int) -> float:
    """Negative log P(token_id) from a logits row (stable log-softmax, numpy)."""
    import numpy as np

    row = np.asarray(logits, dtype=np.float32)
    top = float(row.max())
    logsumexp = top + float(np.log(np.exp(row - top).sum()))
    return logsumexp - float(row[token_id])


def prompt_token_nll(context: str, line: str,
                     max_steps: int = 80) -> list[tuple[str, float]]:
    """Per-token negative log-likelihood of `line`, given a normal `context`.

    One forward pass over `context + line` (logits_all=True), then read the
    conditional logprob of each actual token: logits at position i predict
    token i+1. We return only the tokens belonging to `line` (the prompt tail)
    as (token_text, nll) for the highlighter.

    NOTE: llama-cpp-python's echo-logprobs are misaligned in this version, so we
    read logits directly instead of trusting create_completion(echo=True).
    """
    llm = get_ppl_model()
    full_ids = llm.tokenize((context + line).encode("utf-8"),
                            add_bos=True, special=False)
    full_ids = full_ids[-PPL_N_CTX:]  # keep the tail if longer than the window
    n = len(full_ids)
    if n < 2:
        return []

    llm.reset()
    llm.eval(full_ids)
    scores = llm.scores  # (n_ctx, n_vocab); logits at row i predict token i+1

    line_ids = llm.tokenize(line.encode("utf-8"), add_bos=False, special=False)
    keep = min(max(1, len(line_ids)), max_steps, n - 1)

    out: list[tuple[str, float]] = []
    for j in range(n - keep, n):
        tok = llm.detokenize([full_ids[j]]).decode("utf-8", "replace")
        if j == 0:
            out.append((tok, FLOOR_NLL))
            continue
        out.append((tok, _nll(scores[j - 1], full_ids[j])))
    return out