#!/usr/bin/env python3
"""Verify the training chat template renders reconstructed multi-turn examples
the same way llama.cpp's --jinja rendering of the production GGUF does.

This is a direct response to a real regression class found while building the
retrained corpus: prepare_lora_corpus.py originally emitted tool results as
`role: "tool"` messages. Empirically, hitting a live llama-server instance
(--jinja, loaded with the currently-deployed merged GGUF) with that shape
showed the template silently drops role="tool" content whenever the
preceding assistant turn used the <tena_call> text wrapper instead of native
tool_calls -- which is exactly the wire format this whole system actually
uses (see llm_client.py::_messages_with_text_tool_instructions, called
unconditionally whenever `tools` is passed). The fix was to reconstruct tool
results as `role: "user", content: f"Result for {name}: {content}"`, matching
_strip_native_tool_messages exactly. This script re-runs that same live
verification in a repeatable way so future template or builder changes don't
silently reintroduce the bug.

Usage:
    python3 check_chat_template_parity.py \
        --gguf /var/www/LORA_TenaOS/hf_release/tenaos-gemma4-lora/tenaos-gemma-4-E4B-it-lora-F16.gguf \
        --llama-server /var/www/TenaOS/third_party/llama.cpp/sm80/bin/llama-server \
        --tokenizer-dir /var/www/LORA_TenaOS/hf_release/tenaos-gemma4-lora/merged_hf \
        --samples 2

Requires: a free GPU slot to load the GGUF briefly, and (optionally) a local
transformers install to independently render the same conversations via
AutoTokenizer.apply_chat_template for a second, HF-side comparison.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))
import prepare_lora_corpus as corpus  # noqa: E402

NORMALIZED_ROOT = Path(__file__).resolve().parent / "artifacts_v2" / "normalized"

SAMPLE_KINDS = [
    "form",
    "report",
    "cds",
    "patient_education",
    "scribe_text_english",
    "scribe_text_amharic",
    "voice_scribe_audio",
]


def load_sample_conversations(samples_per_kind: int) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for kind in SAMPLE_KINDS:
        path = NORMALIZED_ROOT / kind / f"{kind}_accepted.jsonl"
        if not path.exists():
            print(f"WARNING: {path} missing, skipping {kind}")
            continue
        count = 0
        with path.open() as handle:
            for line in handle:
                if count >= samples_per_kind:
                    break
                record = json.loads(line)
                examples = corpus.multiturn_examples_from_trace(kind, record)
                for example in examples:
                    out.append(example)
                count += 1
    return out


def check_role_alternation(example: dict[str, Any]) -> list[str]:
    problems = []
    roles = [turn["role"] for turn in example["conversations"]]
    # Tool-calling kinds now carry an extra leading system message (the
    # ephemeral tool-schema instruction llm_client.py always prepends -- see
    # _tool_schema_system_message in prepare_lora_corpus.py), so the prefix
    # is [system, system, user] instead of [system, user]. Scribe kinds never
    # call llm.chat(..., tools=...), so they keep the single-system prefix.
    if roles[:1] == ["system"] and roles[1:2] == ["system"]:
        prefix_len = 3
        if roles[:3] != ["system", "system", "user"]:
            problems.append(f"expected [system, system, user] prefix, got {roles[:3]}")
    else:
        prefix_len = 2
        if roles[:2] != ["system", "user"]:
            problems.append(f"expected [system, user] prefix, got {roles[:2]}")
    rest = roles[prefix_len:]
    # Not strictly alternating any more: cds/patient_education can now emit
    # two consecutive "user" turns (the tool result, then a synthetic
    # steering nudge) -- this matches production exactly, since
    # tool_loop.py/material_loop.py append both as separate messages before
    # the next llm.chat call. Only assert the invariants that always hold:
    # every turn is assistant or user (no "system"/"tool" outside the
    # prefix), and it never starts two assistant turns back-to-back (that
    # would mean a turn boundary was lost during reconstruction).
    prev = None
    for i, role in enumerate(rest):
        if role not in ("assistant", "user"):
            problems.append(f"turn {i + prefix_len}: unexpected role={role!r} outside the system prefix")
            break
        if role == "assistant" and prev == "assistant":
            problems.append(f"turn {i + prefix_len}: two consecutive assistant turns")
            break
        prev = role
    if any(role == "tool" for role in roles):
        problems.append("contains role='tool' — production never sends this role over the wire (see module docstring)")
    return problems


def start_llama_server(llama_server: Path, gguf: Path, port: int) -> subprocess.Popen:
    env_ld = str(llama_server.parent)
    proc = subprocess.Popen(
        [str(llama_server), "-m", str(gguf), "--port", str(port), "--jinja", "-ngl", "999", "-c", "32768"],
        cwd=str(llama_server.parent),
        env={"LD_LIBRARY_PATH": env_ld, "PATH": "/usr/bin:/bin"},
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    return proc


def wait_for_server(port: int, *, timeout_s: float = 180.0) -> None:
    deadline = time.monotonic() + timeout_s
    url = f"http://127.0.0.1:{port}/health"
    while time.monotonic() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=2) as resp:
                if resp.status == 200:
                    return
        except Exception:
            time.sleep(2)
    raise TimeoutError(f"llama-server did not become healthy within {timeout_s}s")


def render_via_llama_server(port: int, conversations: list[dict[str, Any]]) -> list[dict[str, Any]]:
    results = []
    for example in conversations:
        body = json.dumps({"messages": example["conversations"], "max_tokens": 1}).encode("utf-8")
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/chat/completions",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                payload = json.loads(resp.read())
            results.append({
                "id": example["id"],
                "kind": example["kind"],
                "ok": True,
                "prompt_tokens": payload.get("usage", {}).get("prompt_tokens"),
            })
        except Exception as exc:
            results.append({"id": example["id"], "kind": example["kind"], "ok": False, "error": str(exc)})
    return results


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gguf", type=Path, required=False)
    parser.add_argument("--llama-server", type=Path, required=False)
    parser.add_argument("--samples", type=int, default=1, help="Examples per kind to check.")
    parser.add_argument("--port", type=int, default=8899)
    parser.add_argument("--skip-server", action="store_true", help="Only run the static role-alternation check.")
    args = parser.parse_args()

    examples = load_sample_conversations(args.samples)
    print(f"Loaded {len(examples)} sample conversations across {len(SAMPLE_KINDS)} kinds.")

    static_failures = 0
    for example in examples:
        problems = check_role_alternation(example)
        if problems:
            static_failures += 1
            print(f"STATIC FAIL {example['kind']}/{example['id']}: {problems}")
    print(f"Static role-alternation check: {len(examples) - static_failures}/{len(examples)} passed.")

    if args.skip_server:
        raise SystemExit(1 if static_failures else 0)

    proc = start_llama_server(args.llama_server, args.gguf, args.port)
    try:
        wait_for_server(args.port)
        results = render_via_llama_server(args.port, examples)
        server_failures = [r for r in results if not r["ok"]]
        for r in server_failures:
            print(f"SERVER FAIL {r['kind']}/{r['id']}: {r.get('error')}")
        print(f"llama.cpp --jinja render check: {len(results) - len(server_failures)}/{len(results)} passed.")
        exit_code = 1 if (static_failures or server_failures) else 0
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=15)
        except subprocess.TimeoutExpired:
            proc.kill()
    raise SystemExit(exit_code)


if __name__ == "__main__":
    main()