#!/usr/bin/env python3
"""Small, dependency-free llama.cpp performance and context probe."""

from __future__ import annotations

import argparse
import json
import pathlib
import time
import urllib.request


def post(base: str, path: str, payload: dict, timeout: int = 1800) -> tuple[dict, float]:
    request = urllib.request.Request(
        base + path,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    started = time.monotonic()
    with urllib.request.urlopen(request, timeout=timeout) as response:
        result = json.load(response)
    return result, time.monotonic() - started


def make_text(lines: int) -> str:
    return "\n".join(
        f"Record {n:06d}: cobalt lantern maple orbit quartz river silver tango." for n in range(lines)
    )


def count_tokens(base: str, text: str) -> int:
    result, _ = post(base, "/tokenize", {"content": text, "add_special": False})
    return len(result.get("tokens", []))


def sized_text(base: str, target: int) -> tuple[str, int]:
    # One probe establishes the tokenizer-specific tokens per synthetic line.
    sample = make_text(100)
    per_line = max(1.0, count_tokens(base, sample) / 100)
    lines = max(1, int(target / per_line))
    text = make_text(lines)
    actual = count_tokens(base, text)
    if actual < target * 0.95:
        lines = int(lines * target / max(1, actual))
        text = make_text(lines)
        actual = count_tokens(base, text)
    return text, actual


def chat(base: str, prompt: str, max_tokens: int, temperature: float = 0.2) -> dict:
    result, wall = post(base, "/v1/chat/completions", {
        "model": "benchmark",
        "temperature": temperature,
        "max_tokens": max_tokens,
        "reasoning_effort": "none",
        "messages": [{"role": "user", "content": prompt}],
    })
    message = (result.get("choices") or [{}])[0].get("message") or {}
    return {
        "wall_seconds": round(wall, 3),
        "timings": result.get("timings", {}),
        "usage": result.get("usage", {}),
        "content": message.get("content", ""),
        "reasoning_content": message.get("reasoning_content", ""),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("label")
    parser.add_argument("context", type=int)
    parser.add_argument("--base", default="http://127.0.0.1:5004")
    parser.add_argument("--output", default="/data/benchmarks/dirk-qwen38")
    args = parser.parse_args()

    result: dict = {"label": args.label, "context": args.context, "started": time.time()}
    short, short_n = sized_text(args.base, min(16000, max(4000, args.context // 10)))
    prompt = short + "\nReply with exactly: PREFILL-OK"
    result["prompt_tokens_synthetic"] = short_n
    result["uncached"] = chat(args.base, prompt, 32)
    result["cached"] = chat(args.base, prompt, 32)

    output_prompt = (
        "Return exactly 256 comma-separated integers beginning at 1 and ending at 256. "
        "Do not explain and do not omit any integer."
    )
    result["decode"] = chat(args.base, output_prompt, 768)

    recall_target = int(args.context * 0.70)
    long_text, long_n = sized_text(args.base, recall_target)
    marks = [
        (len(long_text) // 8, "NEEDLE_ALPHA=RAVEN-417"),
        (len(long_text) // 2, "NEEDLE_BETA=CEDAR-928"),
        (len(long_text) * 7 // 8, "NEEDLE_GAMMA=ORBIT-563"),
    ]
    for position, needle in reversed(marks):
        long_text = long_text[:position] + "\n" + needle + "\n" + long_text[position:]
    recall_prompt = long_text + (
        "\nReturn only a JSON object with keys alpha, beta, gamma and their exact values "
        "from the three NEEDLE lines."
    )
    recall = chat(args.base, recall_prompt, 256)
    recall["synthetic_tokens"] = long_n
    content = recall.get("content", "")
    recall["needles_found"] = {
        "alpha": "RAVEN-417" in content,
        "beta": "CEDAR-928" in content,
        "gamma": "ORBIT-563" in content,
    }
    result["recall"] = recall
    result["finished"] = time.time()

    output = pathlib.Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    target = output / f"{args.label}.json"
    target.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n")
    print(json.dumps(result, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())