"""Local GGUF LLM client (llama.cpp). Lazy-loaded + cached; no network at generate time.

``llama_cpp`` is imported INSIDE the loader (not at module top) so importing this module — and running the
gates — needs no compiled backend or downloaded weights. Models cache under ``data/artifacts/models/`` via
an ``HF_HOME`` redirect, so once fetched the run is fully offline.
"""

from __future__ import annotations

import os
import re
import threading
import time
from functools import lru_cache
from typing import TYPE_CHECKING, Any

from common.logging import get_logger
from common.paths import REPO_ROOT

if TYPE_CHECKING:
    from config import Settings

MODEL_DIR = REPO_ROOT / "data" / "artifacts" / "models"


def cached_hf_download(repo_id: str, filename: str) -> str:
    """Resolve ``filename`` from ``repo_id``: local cache first (zero network), download only if missing.

    Plain ``hf_hub_download`` fires a live HEAD request on every call — even for a fully-cached file — to
    validate the etag, which is a real network attempt the "no network during ranking" contract shouldn't
    have to tolerate. ``try_to_load_from_cache`` gives a tri-state local-only answer (found / confirmed
    absent via a ``.no_exist`` marker from a prior check / never checked) — only the last state should ever
    reach the network, otherwise a known-missing optional file (e.g. a model without a sibling
    ``.onnx_data``) would retry the Hub on every single load forever."""
    from huggingface_hub import _CACHED_NO_EXIST, hf_hub_download, try_to_load_from_cache

    cached = try_to_load_from_cache(repo_id, filename)
    if isinstance(cached, str):
        return cached
    if cached is _CACHED_NO_EXIST:
        raise FileNotFoundError(f"{filename!r} confirmed absent from {repo_id!r} (cached .no_exist marker)")
    return hf_hub_download(repo_id=repo_id, filename=filename)  # never checked locally — genuinely need network

# Qwen3 is a reasoning model — it emits <think>…</think> before its answer. We never want that prose in the
# extracted spec / HyDE, so we both ask it to skip thinking (/no_think soft switch) and strip any residual
# think block (incl. a truncated, unclosed one) from the output.
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
_THINK_DANGLING = re.compile(r"<think>.*$", re.DOTALL)


def _strip_thinking(text: str) -> str:
    """Remove a Qwen3 <think> block (closed or truncated) and surrounding whitespace."""
    return _THINK_DANGLING.sub("", _THINK_BLOCK.sub("", text)).strip()


_HEARTBEAT_SECONDS = 5.0


def _generate_streaming(model: Any, messages: list, temperature: float, top_p: float, max_tokens: int, log: Any) -> str:
    """Run ``create_chat_completion(stream=True)`` and log a heartbeat every ``_HEARTBEAT_SECONDS`` — gives
    visibility into whether a long call is actively progressing (chars accumulating) or genuinely stuck
    (heartbeats stop appearing), which a single before/after log pair can't distinguish. Used for EVERY call,
    not just timed ones — JD parse has no timeout by design, so this is the only way to tell "slow but working"
    from "hung" without cutting anything off."""
    chunks: list[str] = []
    start = time.perf_counter()
    last_heartbeat = start
    for chunk in model.create_chat_completion(
        messages=messages, temperature=temperature, top_p=top_p, max_tokens=max_tokens, stream=True
    ):
        piece = chunk["choices"][0].get("delta", {}).get("content") or ""
        if piece:
            chunks.append(piece)
        now = time.perf_counter()
        if now - last_heartbeat >= _HEARTBEAT_SECONDS:
            log.info("llm.generate_heartbeat", elapsed=round(now - start, 1), chars_so_far=sum(len(c) for c in chunks))
            last_heartbeat = now
    return "".join(chunks)


def _run_with_timeout(fn: Any, timeout: float) -> Any:
    """Run ``fn()`` on a daemon thread and wait up to ``timeout`` seconds; raise ``TimeoutError`` if it's still
    running. llama.cpp offers no cooperative cancellation mid-generation, so a genuinely stuck/runaway call
    can't be killed — the thread is simply abandoned (daemon=True so it can't block process exit) and the
    caller moves on. This exists because a single generation was observed taking 1996s (33 min) on real
    deployment hardware with NOTHING to stop it — the between-call budget check in tiers.py can't help since
    it only gates the NEXT call, not one already in flight."""
    result: dict[str, Any] = {}

    def target() -> None:
        try:
            result["value"] = fn()
        except Exception as exc:  # noqa: BLE001 — re-raised on the calling thread below
            result["error"] = exc

    thread = threading.Thread(target=target, daemon=True)
    thread.start()
    thread.join(timeout)
    if thread.is_alive():
        raise TimeoutError(f"generation exceeded {timeout}s — abandoned (still running in the background)")
    if "error" in result:
        raise result["error"]
    return result["value"]


@lru_cache(maxsize=2)
def _load_llama(repo: str, gguf: str, n_ctx: int, n_threads: int) -> Any:
    """Load a GGUF model into a cached llama.cpp handle.

    An EXACT ``gguf`` filename is resolved with ``cached_hf_download`` (cache-only lookup first, no Hub API
    call once baked — falls back to a real fetch only if missing). ``Llama.from_pretrained`` is NOT used for
    exact names because it always lists the repo over the network. A glob ``gguf`` (the rare fallback model)
    still routes through ``from_pretrained`` (needs network)."""
    from llama_cpp import Llama  # imported lazily — keeps the import graph backend-free

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(MODEL_DIR))  # cache weights inside the repo, not the user home
    log = get_logger("models")
    log.info("llm.load", repo=repo, gguf=gguf, n_ctx=n_ctx, n_threads=n_threads)
    start = time.perf_counter()
    try:
        if "*" in gguf or "?" in gguf:  # glob → must list the repo (online only); used by the SmolLM2 fallback
            model = Llama.from_pretrained(repo_id=repo, filename=gguf, n_ctx=n_ctx,
                                          n_threads=(n_threads or None), verbose=False)
        else:
            model_path = cached_hf_download(repo, gguf)  # local cache first; downloads only if genuinely missing
            model = Llama(model_path=model_path, n_ctx=n_ctx, n_threads=(n_threads or None), verbose=False)
    except Exception as exc:
        log.error("llm.load_failed", repo=repo, gguf=gguf, seconds=round(time.perf_counter() - start, 3), error=str(exc))
        raise
    log.info("llm.load_complete", repo=repo, seconds=round(time.perf_counter() - start, 3))
    return model


class LlmClient:
    """Thin wrapper over a cached GGUF model exposing a single ``generate`` call."""

    def __init__(self, repo: str, gguf: str, n_ctx: int = 4096, n_threads: int = 0) -> None:
        self._repo, self._gguf, self._n_ctx, self._n_threads = repo, gguf, n_ctx, n_threads
        # ``/no_think`` is a Qwen3-only soft switch; appending it to other models (LFM2.5, etc.) would inject
        # literal junk into the prompt. Gate it on the repo so each model gets the right convention.
        self._soft_no_think = "qwen" in repo.lower()

    def ensure(self) -> None:
        """Force the (down)load now — used by setup so the rank/JD run doesn't pay it on the clock."""
        _load_llama(self._repo, self._gguf, self._n_ctx, self._n_threads)

    def generate(
        self, prompt: str, system: str | None = None,
        temperature: float = 0.0, top_p: float = 0.95, max_tokens: int = 512, think: bool = False,
        timeout: float | None = None,
    ) -> str:
        """Chat-complete ``prompt`` (temperature=0 ⇒ greedy). Reasoning suppressed + stripped by default.

        Always streams internally with a periodic heartbeat log (``llm.generate_heartbeat`` every ~5s) — for a
        call with no ``timeout`` (JD parse, by design), that heartbeat is the ONLY way to tell "slow but still
        producing tokens" from "genuinely hung" during a long silence; without it there's nothing between
        ``generate_start`` and either ``generate_complete`` or nothing at all.

        ``timeout``, if given, is a hard wall-clock cap enforced via ``_run_with_timeout`` — a real generation
        was observed taking 1996s with nothing to stop it, so callers on a wall-clock budget (Step 11 polish)
        should always pass one; JD parse intentionally does NOT (its accuracy matters more than its wall-clock).
        """
        model = _load_llama(self._repo, self._gguf, self._n_ctx, self._n_threads)
        # Qwen3: append the /no_think soft switch to skip chain-of-thought. Non-Qwen models: send as-is.
        user = prompt + " /no_think" if self._soft_no_think and not think else prompt
        messages = ([{"role": "system", "content": system}] if system else []) + [
            {"role": "user", "content": user}
        ]
        log = get_logger("models")
        log.info("llm.generate_start", prompt_chars=len(user), max_tokens=max_tokens, timeout=timeout)
        start = time.perf_counter()
        call = lambda: _generate_streaming(model, messages, temperature, top_p, max_tokens, log)  # noqa: E731
        try:
            raw_text = _run_with_timeout(call, timeout) if timeout else call()
        except Exception as exc:
            log.error("llm.generate_failed", seconds=round(time.perf_counter() - start, 3), error=str(exc))
            raise
        text = _strip_thinking(raw_text)
        log.info("llm.generate_complete", seconds=round(time.perf_counter() - start, 3), chars_out=len(text))
        return text


def index_gen_client(settings: Settings) -> LlmClient:
    """The JD-parse LLM client (Qwen3-1.7B GGUF, standard instruct) built from central config.

    ``models.n_threads`` is llama.cpp's OWN explicit override; when it's left at its default (0), the shared
    ``compute.cpu_threads`` budget applies instead, so one config knob caps the whole engine consistently.
    """
    from common.runtime import resolve_cpu_threads

    models = settings.models
    threads = models.n_threads or resolve_cpu_threads(settings.compute.cpu_threads)
    return LlmClient(models.index_gen_llm_repo, models.index_gen_llm_gguf, models.n_ctx, threads)