"""Run one chat turn against the Q4_K_M GGUF of DeepSeek-R1-Distill-Qwen-1.5B with llama.cpp."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from llama_cpp import Llama


HERE = Path(__file__).resolve().parent
MODEL_FILENAME = "deepseek-ai__DeepSeek-R1-Distill-Qwen-1.5B_llamacpp_Q4_K_M.gguf"

DEFAULT_PROMPT = (
    "What is the smallest positive integer that is divisible by every integer "
    "from 1 to 10? Put your final answer in \\boxed{}."
)

# The sampling settings the published MATH-500 pass@1 was measured under.
TEMPERATURE = 0.6
TOP_P = 0.95
MAX_TOKENS = 12288
# N_CTX is the shipped serving window, not the evaluation one: MATH-500 ran at
# 14336 so the whole 12288-token budget fits. Raise this to reproduce that
# contract — at 4096 a long trace hits the context wall before max_tokens.
N_CTX = 4096
N_THREADS = 4

THINK_CLOSE = "</think>"


def resolve_model_path(explicit: str | None) -> Path:
    """Find the GGUF next to this script, or one directory up."""
    if explicit:
        path = Path(explicit).expanduser().resolve()
        if not path.is_file():
            raise FileNotFoundError(f"No GGUF at {path}")
        return path

    for candidate in (HERE / MODEL_FILENAME, HERE.parent / MODEL_FILENAME):
        if candidate.is_file():
            return candidate

    raise FileNotFoundError(
        f"{MODEL_FILENAME} not found in {HERE} or {HERE.parent}. "
        "Pass --model with an explicit path."
    )


def split_reasoning(text: str) -> tuple[str, str]:
    """Separate the chain of thought from the final answer.

    Returns ``(reasoning, answer)``. When the trace was cut off before
    ``</think>`` the whole output is reasoning and the answer is empty — that is
    a truncation, not a refusal, and it means the token budget was too small.
    """
    if THINK_CLOSE in text:
        reasoning, answer = text.split(THINK_CLOSE, 1)
        return reasoning.strip(), answer.strip()
    return text.strip(), ""


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default=None, help="Path to the .gguf file")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--max-tokens", type=int, default=MAX_TOKENS)
    parser.add_argument("--temperature", type=float, default=TEMPERATURE)
    parser.add_argument("--top-p", type=float, default=TOP_P)
    parser.add_argument("--n-ctx", type=int, default=N_CTX)
    parser.add_argument("--threads", type=int, default=N_THREADS)
    args = parser.parse_args()

    model_path = resolve_model_path(args.model)
    print(f"Loading {model_path.name} ({model_path.stat().st_size / 1024**2:.2f} MB)")

    llm = Llama(
        model_path=str(model_path),
        n_ctx=args.n_ctx,
        n_threads=args.threads,
        verbose=False,
    )

    # create_chat_completion applies the chat template stored in the GGUF, which
    # is what appends the opening <think> tag that puts the model in reasoning
    # mode. Hand-building the prompt string without it changes the behaviour.
    completion = llm.create_chat_completion(
        messages=[{"role": "user", "content": args.prompt}],
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
    )

    text = completion["choices"][0]["message"]["content"]
    finish_reason = completion["choices"][0]["finish_reason"]
    reasoning, answer = split_reasoning(text)

    print(f"\nPrompt: {args.prompt}")
    print(f"\nReasoning ({len(reasoning)} chars, truncated below):\n{reasoning[:800]}")
    print(f"\nFinal answer:\n{answer or '(none — the trace was truncated)'}")
    print(f"\nfinish_reason: {finish_reason}")
    print(f"tokens: {completion['usage']}")

    output_path = HERE / "generation.json"
    output_path.write_text(
        json.dumps(
            {
                "model": model_path.name,
                "prompt": args.prompt,
                "reasoning": reasoning,
                "answer": answer,
                "finish_reason": finish_reason,
                "usage": completion["usage"],
                "sampling": {
                    "temperature": args.temperature,
                    "top_p": args.top_p,
                    "max_tokens": args.max_tokens,
                    "n_ctx": args.n_ctx,
                    "n_threads": args.threads,
                },
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(f"\nWrote {output_path}")


if __name__ == "__main__":
    main()