"""Title models in Ollama on eval/messages.jsonl, sent the way Simplixity sends them.

The request is the one app/llm/client.py builds for a title on Ollama: the
OpenAI-compatible endpoint, temperature 0.2, max_tokens 24 and
reasoning_effort "none" (reasoning_tokens=0 adds no headroom on top). The
system prompt is read from the training data, which is the app's title prompt
verbatim (TITLE and IN_ENGLISH from app/agent/prompts.py).

    python3 -I eval.py simplixity-title:0.8b-v1-q8_0 simplixity-title:0.8b-v2-q8_0

A model given as `label=hf:/path/to/checkpoint` runs in transformers on the GPU
at bf16 instead, for checkpoints this Ollama cannot import: 0.40 on Linux takes
safetensors only through its MLX runtime, which the Arch package leaves out.
It renders the same chat template (the generation prompt ends in the closed
empty thinking block, as with reasoning_effort "none") and samples like
Ollama's defaults: temperature 0.2, top_k 40, top_p 0.9.

    uv run python -I eval/eval.py v2=hf:runs/v2/model simplixity-title:0.8b-v1-q8_0
"""

from __future__ import annotations

import argparse
import json
import os
import re
import statistics
import time
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
TRAIN = os.path.expanduser("~/models/simplixity-title-0.8b/training/train.jsonl")
TOKENIZER = os.path.expanduser("~/models/simplixity-title-0.8b")

QUOTES = set("\"“”„«»‚‘`")
# A title is German if it has one of these, a word with an umlaut or ß, or a
# word of five letters or more copied from a German message ("Kurzarbeit
# Gehalt Bedeutung" has neither of the first two). Names the keys list
# (Großglockner, BahnCard, Mietpreisbremse) and the words below, which are
# English as well, do not count.
SHARED_WORDS = {"restaurant", "docker", "compose", "python", "commit", "laptop", "blockchain",
                "filesharing", "gravelbike", "ibuprofen", "tirol", "pandas", "semikolon", "winter"}
GERMAN_WORDS = {
    "der", "die", "das", "und", "mit", "für", "von", "im", "ist", "ein", "eine",
    "zu", "auf", "bei", "wie", "nicht", "oder", "den", "dem", "des", "zum", "zur",
    "über", "nach", "vom", "am", "beim", "unterschied", "frage", "tipps",
}


def system_prompt() -> str:
    with open(TRAIN) as f:
        return json.loads(f.readline())["messages"][0]["content"]


def ask(base: str, model: str, system: str, text: str) -> tuple[str, float]:
    payload = {
        "model": model,
        "messages": [{"role": "system", "content": system}, {"role": "user", "content": text[:1000]}],
        "temperature": 0.2,
        "stream": False,
        "max_tokens": 24,
        "reasoning_effort": "none",
    }
    req = urllib.request.Request(
        f"{base}/v1/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.perf_counter()
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.load(resp)
    elapsed = time.perf_counter() - start
    return (data["choices"][0]["message"].get("content") or "").strip(), elapsed


class Local:
    def __init__(self, path: str) -> None:
        import torch
        from transformers import AutoModelForImageTextToText, AutoTokenizer

        self.torch = torch
        config = os.path.join(path, "adapter_config.json")
        if os.path.exists(config):
            # A training checkpoint: LoRA on top of the model it was trained from.
            from peft import PeftModel

            with open(config) as f:
                init = json.load(f)["base_model_name_or_path"]
            self.tok = AutoTokenizer.from_pretrained(TOKENIZER)
            base = AutoModelForImageTextToText.from_pretrained(init, dtype=torch.bfloat16)
            self.model = PeftModel.from_pretrained(base, path).merge_and_unload().cuda().eval()
        else:
            self.tok = AutoTokenizer.from_pretrained(path)
            self.model = AutoModelForImageTextToText.from_pretrained(path, dtype=torch.bfloat16).cuda().eval()
        # The checkpoints have no generation_config.json, so generate would stop
        # only at config.json's <|endoftext|>. Ollama stops at the end of the
        # turn, <|im_end|>; without it v1 keeps writing "\n!\n!..." to 24 tokens.
        self.stop = [self.tok.convert_tokens_to_ids("<|im_end|>"), self.tok.eos_token_id]

    def ask(self, system: str, text: str) -> tuple[str, float]:
        messages = [{"role": "system", "content": system}, {"role": "user", "content": text[:1000]}]
        prompt = self.tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
        ids = self.tok(prompt, add_special_tokens=False, return_tensors="pt").to("cuda")
        start = time.perf_counter()
        with self.torch.no_grad():
            out = self.model.generate(
                **ids, max_new_tokens=24, do_sample=True, temperature=0.2, top_k=40, top_p=0.9,
                eos_token_id=self.stop,
            )
        elapsed = time.perf_counter() - start
        title = self.tok.decode(out[0, ids["input_ids"].shape[1] :], skip_special_tokens=True)
        return title.strip(), elapsed


def judge(title: str, msg: dict) -> dict:
    words = title.split()
    lower = title.lower()
    tokens = {w.strip(".,:;!?()'’").lower() for w in words}

    def name(w: str) -> bool:
        return w in SHARED_WORDS or any(k in w or w in k for k in msg["keys"])

    source = set(re.findall(r"\w+", msg["text"].lower())) if msg["lang"] == "de" else set()
    umlaut = any(re.search(r"[äöüß]", w) and not name(w) for w in tokens)
    copied = any(len(w) >= 5 and w in source and not name(w) for w in tokens)
    checks = {
        "empty": not title,
        "too_long": len(words) > 6,
        "quotes": any(c in QUOTES for c in title) or title[:1] == "'" or title[-1:] == "'",
        "trailing_punct": title[-1:] in set(".!?:;,…"),
        "chat_about": "chat about" in lower,
        "german": bool(tokens & GERMAN_WORDS) or umlaut or copied,
    }
    return {
        "format_ok": not any(checks.values()),
        "fails": [k for k, v in checks.items() if v],
        "key_ok": any(k in lower for k in msg["keys"]) if msg["keys"] else None,
    }


def pct(xs: list[bool]) -> str:
    return f"{100 * sum(xs) / len(xs):.0f}%" if xs else "–"


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("models", nargs="*")
    p.add_argument("--runs", type=int, default=3)
    p.add_argument("--base", default="http://localhost:11434")
    p.add_argument("--out", default=os.path.join(HERE, "results.jsonl"))
    p.add_argument("--messages", default=os.path.join(HERE, "messages.jsonl"))
    p.add_argument("--rejudge", action="store_true", help="score the titles in --out again, generate nothing")
    args = p.parse_args()

    system = system_prompt()
    with open(args.messages) as f:
        messages = [json.loads(line) for line in f]

    rows = []
    labels = []
    if args.rejudge:
        by_id = {m["id"]: m for m in messages}
        with open(args.out) as f:
            rows = [json.loads(line) for line in f]
        rows = [{**r, **judge(r["title"], by_id[r["id"]])} for r in rows]
        labels = list(dict.fromkeys(r["model"] for r in rows))
        args.runs = len({r["run"] for r in rows})
        args.models = []
    for spec in args.models:
        model, _, path = spec.partition("=hf:")
        labels.append(model)
        if path:
            local = Local(os.path.expanduser(path))
            call = local.ask
        else:
            def call(system: str, text: str, model: str = model) -> tuple[str, float]:
                return ask(args.base, model, system, text)
        call(system, "hello")  # load it before anything is timed
        for run in range(args.runs):
            for msg in messages:
                title, elapsed = call(system, msg["text"])
                rows.append({"model": model, "run": run, "id": msg["id"], "lang": msg["lang"],
                             "title": title, "seconds": elapsed, **judge(title, msg)})
            print(f"{model} run {run + 1}/{args.runs} done", flush=True)
        if path:
            del local, call
            import torch

            torch.cuda.empty_cache()
    args.models = labels
    with open(args.out, "w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"\n{len(messages)} messages × {args.runs} runs\n")
    head = ["", *args.models]
    lines = [head, ["---"] * len(head)]
    for label, sel in (("all", None), ("en", "en"), ("de", "de")):
        cell = []
        for model in args.models:
            rs = [r for r in rows if r["model"] == model and sel in (None, r["lang"])]
            cell.append(pct([r["format_ok"] for r in rs]))
        lines.append([f"All format rules ({label})", *cell])
    for fail in ("too_long", "german", "quotes", "trailing_punct", "empty", "chat_about"):
        lines.append([f"  fails: {fail}", *[
            pct([fail in r["fails"] for r in rows if r["model"] == m]) for m in args.models]])
    lines.append(["Key term named", *[
        pct([r["key_ok"] for r in rows if r["model"] == m and r["key_ok"] is not None]) for m in args.models]])
    lines.append(["Median per title", *[
        f"{statistics.median(r['seconds'] for r in rows if r['model'] == m):.2f} s" for m in args.models]])
    for line in lines:
        print("| " + " | ".join(line) + " |")


if __name__ == "__main__":
    main()