"""Validate the Llama-3 tool template against the real corpus, off-GPU.

Everything here runs on jinja2 + tokenizers alone -- no torch, no VM -- so the
template is proven before any GPU time is spent. It checks the two things that
fail silently rather than loudly:

  * tool calls and tool schemas survive rendering (the stock Sahabat-AI
    template drops both, and an assistant tool-call turn has `content: ""`, so
    the damage looks like an ordinary empty reply);
  * response-only masking covers exactly the assistant turns. A tool result
    rendered under a header the instruction delimiter does not match would sit
    inside the loss and teach the model to write its own forecasts.

    python3 training/check_template.py --data train.parquet \
        --tokenizer tokenizer.json
"""

import argparse
import json
import statistics
from pathlib import Path

import jinja2.ext
import pyarrow.parquet as pq
from jinja2.sandbox import ImmutableSandboxedEnvironment

TEMPLATE = Path(__file__).resolve().parent / "llama3_tools.jinja"

BOS = "<|begin_of_text|>"
INSTRUCTION_PART = "<|start_header_id|>user<|end_header_id|>\n\n"
RESPONSE_PART = "<|start_header_id|>assistant<|end_header_id|>\n\n"


def build_env():
    """Reproduce transformers' template environment closely enough to trust."""
    env = ImmutableSandboxedEnvironment(
        trim_blocks=True, lstrip_blocks=True, extensions=[jinja2.ext.loopcontrols])
    env.filters["tojson"] = lambda x, **kw: json.dumps(
        x, ensure_ascii=kw.get("ensure_ascii", False), indent=kw.get("indent"))
    env.globals["raise_exception"] = lambda m: (_ for _ in ()).throw(Exception(m))
    return env


def hydrate(row):
    messages = []
    for message in row["messages"]:
        clean = {k: v for k, v in message.items() if v is not None}
        if "tool_calls" in clean:
            clean["tool_calls"] = json.loads(clean["tool_calls"])
        messages.append(clean)
    tools = row.get("tools")
    return messages, (json.loads(tools) if tools else None)


def trained_spans(text):
    """The character spans train_on_responses_only would keep in the loss.

    unsloth trains from each response delimiter to the next *instruction*
    delimiter, or to the end. Working in characters rather than token ids is
    equivalent here and keeps the check tokenizer-free.
    """
    spans, cursor = [], 0
    while (start := text.find(RESPONSE_PART, cursor)) != -1:
        start += len(RESPONSE_PART)
        stop = text.find(INSTRUCTION_PART, start)
        stop = len(text) if stop == -1 else stop
        spans.append((start, stop))
        cursor = stop
    return spans


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True,
                        help="train.parquet from prepare_data.py")
    parser.add_argument("--tokenizer", help="path to a tokenizer.json, for length stats")
    parser.add_argument("--rows", type=int, default=4000)
    args = parser.parse_args()

    template = build_env().from_string(TEMPLATE.read_text(encoding="utf-8"))
    rows = pq.read_table(args.data).slice(0, args.rows).to_pylist()

    calls = schemas = leaked = empty_mask = 0
    tool_rows = 0
    lengths = []
    tokenizer = None
    if args.tokenizer:
        from tokenizers import Tokenizer
        tokenizer = Tokenizer.from_file(args.tokenizer)

    for row in rows:
        messages, tools = hydrate(row)
        text = template.render(messages=messages, tools=tools, bos_token=BOS,
                               add_generation_prompt=False)

        spans = trained_spans(text)
        if not spans:
            empty_mask += 1
        trained = "".join(text[a:b] for a, b in spans)

        # A tool result inside the trained span is the failure this exists for.
        leaked += trained.count("<tool_response>")

        for message in messages:
            for call in message.get("tool_calls") or []:
                tool_rows += 1
                needle = f'{{"name": "{call["function"]["name"]}", "parameters": '
                calls += needle in text
                # ...and the call itself MUST be trained on, or the model never
                # learns to emit one.
                calls -= needle not in trained
        if tools:
            schemas += all(t["function"]["description"][:40] in text for t in tools)

        if tokenizer:
            lengths.append(len(tokenizer.encode(text, add_special_tokens=False).ids))

    print(f"rows                 {len(rows)}")
    print(f"tool calls           {calls}/{tool_rows} rendered and inside the loss")
    print(f"tool schemas         {schemas} conversations carry the offered schema")
    print(f"tool results leaked  {leaked}  (must be 0)")
    print(f"empty label masks    {empty_mask}  (must be 0)")

    if lengths:
        lengths.sort()
        n = len(lengths)
        print(f"tokens               mean={statistics.mean(lengths):.0f} "
              f"p50={lengths[n // 2]} p90={lengths[int(n * .9)]} "
              f"p99={lengths[int(n * .99)]} max={lengths[-1]} "
              f"over_4096={sum(1 for x in lengths if x > 4096)}")

    if leaked or empty_mask or calls != tool_rows:
        raise SystemExit("template check FAILED")
    print("template check OK")


if __name__ == "__main__":
    main()