#!/usr/bin/env python3
"""Is the Sharp template actually applied to this model? Check, don't guess.

    python3 scripts/check_applied.py /path/to/model-dir
    python3 scripts/check_applied.py model.gguf

A model directory can carry the chat template in TWO places -- `chat_template.jinja` and the
`chat_template` key inside `tokenizer_config.json` -- and runtimes disagree about which one
wins. Recent transformers prefers the .jinja file; oMLX and several others read the embedded
copy and ignore the file entirely. So dropping in a new .jinja can appear to do nothing, with
no error anywhere, and the only symptom is that the model still writes preamble.

This reports every source it finds, renders each one, and tells you whether they agree. If they
disagree, it says so loudly -- that mismatch is the actual bug, and it is invisible otherwise.

Exit code is 0 only if every template source present has the terseness prompt applied.
"""
from __future__ import annotations

import json
import pathlib
import sys

MARKER = "Never: open with preamble"
SYSTEM_PROBE = "Be a pirate."


def render(src: str, msgs: list[dict], **kw) -> str:
    try:
        from jinja2 import Environment
    except ImportError:
        sys.exit("needs jinja2:  pip install jinja2")
    return Environment().from_string(src).render(
        messages=msgs, add_generation_prompt=True, **kw)


def describe(src: str) -> dict:
    """Render the cases that matter and report what the template does."""
    user = [{"role": "user", "content": "hi"}]
    with_sys = [{"role": "system", "content": SYSTEM_PROBE}, {"role": "user", "content": "hi"}]
    multi = [{"role": "user", "content": "Q1"},
             {"role": "assistant", "content": "<think>t</think>A1"},
             {"role": "user", "content": "Q2"}]
    try:
        plain, sysd, mt = render(src, user), render(src, with_sys), render(src, multi)
    except Exception as e:                       # a template that won't render is its own answer
        return {"error": f"{type(e).__name__}: {e}"}
    return {
        "terse_count": plain.count(MARKER),
        "keeps_system": SYSTEM_PROBE in sysd,
        "retains_think": "<think>t</think>" in mt,
        "identity": next((n for n in ("Nail-35b-a3b", "Dagger-27b") if n in plain), None),
        "bytes": len(src),
    }


def report(label: str, src: str) -> bool:
    d = describe(src)
    print(f"\n  [{label}]  {d.get('bytes', 0)} bytes")
    if "error" in d:
        print(f"     FAILS TO RENDER — {d['error']}")
        return False
    n = d["terse_count"]
    ok = n == 1
    print(f"     terseness prompt ......... {'yes' if n == 1 else f'NO (found {n}x)'}")
    print(f"     keeps your system prompt . {'yes' if d['keeps_system'] else 'NO'}")
    # Heuristic: a template that simply echoes message content will "pass" this without
    # implementing retention at all. Reliable as a NO, only suggestive as a yes.
    print(f"     retains thinking* ........ {'yes' if d['retains_think'] else 'no'}"
          f"{'' if d['retains_think'] else '  (stock behaviour, not froggeric-fixed)'}")
    if d["identity"]:
        print(f"     WARNING: names a specific model ({d['identity']}) — you probably want the")
        print(f"              model-agnostic template from this repo instead")
    return ok


def from_gguf(path: pathlib.Path) -> str | None:
    try:
        from gguf import GGUFReader
    except ImportError:
        sys.exit("reading a .gguf needs the gguf package:  pip install gguf")
    r = GGUFReader(str(path))
    f = r.fields.get("tokenizer.chat_template")
    if f is None:
        return None
    v = f.contents()
    return v if isinstance(v, str) else None


def main() -> int:
    if len(sys.argv) != 2:
        sys.exit(__doc__.strip().splitlines()[2].strip())
    target = pathlib.Path(sys.argv[1]).expanduser()
    if not target.exists():
        sys.exit(f"no such path: {target}")

    print(f"checking {target}")
    sources: dict[str, str] = {}

    if target.is_file() and target.suffix == ".gguf":
        tpl = from_gguf(target)
        if tpl is None:
            print("\n  no tokenizer.chat_template embedded — llama.cpp will fall back to a\n"
                  "  built-in template, and the terseness prompt is NOT applied.")
            return 1
        sources["embedded in .gguf"] = tpl
    else:
        j = target / "chat_template.jinja"
        if j.is_file():
            sources["chat_template.jinja"] = j.read_text()
        tc = target / "tokenizer_config.json"
        if tc.is_file():
            key = json.loads(tc.read_text()).get("chat_template")
            if isinstance(key, str):
                sources["tokenizer_config.json"] = key
        if not sources:
            print("\n  no chat template found at all — nothing is applied.")
            return 1

    ok = all([report(name, src) for name, src in sources.items()])

    if len(sources) > 1:
        vals = list(sources.values())
        agree = all(v == vals[0] for v in vals)
        print()
        if agree:
            print("  Both sources carry the SAME template — whichever your runtime prefers,")
            print("  you get the same behaviour.")
        else:
            ok = False
            print("  *** THE TWO SOURCES DISAGREE ***")
            print("  Recent transformers uses chat_template.jinja; oMLX and others read the")
            print("  copy embedded in tokenizer_config.json. Right now those differ, so what")
            print("  you get depends on your runtime. Patch both to the same content.")

    print("\n  * retention is inferred from the rendered output; a template that merely echoes"
          "\n    message content passes it without implementing retention. Trust the 'no'.")
    print("\n" + ("APPLIED" if ok else "NOT APPLIED (or inconsistent) — see above"))
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())