"""End-to-end test: constrained decoding for Gemma-4 tool calling with the modified xgrammar.

Loads google/gemma-4-E2B-it via HF transformers, builds the gemma_4 builtin
structural tag from xgrammar, and generates with xgrammar's LogitsProcessor.
Verifies:
  1. The grammar compiles against the real Gemma-4 tokenizer.
  2. Constrained generation emits a well-formed <|tool_call>call:name{...}<tool_call|>.
  3. The emitted arguments use Gemma's <|"|> string delimiters and satisfy the schema.
  4. (Alignment) Unconstrained generation is compared against the same grammar.
"""

import json
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import xgrammar as xgr
from xgrammar.testing import _is_grammar_accept_string

MODEL_ID = "google/gemma-4-E2B-it"

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "City name"},
                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                },
                "required": ["city"],
            },
        },
    }
]

MESSAGES = [{"role": "user", "content": "What's the weather in Seoul in celsius?"}]


def main():
    print(f"=== loading tokenizer/model: {MODEL_ID} ===")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, device_map="mps")
    model.eval()

    vocab_size = model.config.get_text_config().vocab_size
    print(f"vocab_size={vocab_size}")

    prompt = tokenizer.apply_chat_template(
        MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False
    )
    print("=== rendered prompt (tail) ===")
    print(prompt[-600:])
    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("mps")

    tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
    compiler = xgr.GrammarCompiler(tokenizer_info)

    results = {}
    for tool_choice, reasoning in [("required", False), ("auto", True)]:
        label = f"tool_choice={tool_choice}, reasoning={reasoning}"
        print(f"\n=== structural tag: {label} ===")
        stag = xgr.get_model_structural_tag(
            model="gemma_4", tools=TOOLS, tool_choice=tool_choice, reasoning=reasoning
        )
        t0 = time.monotonic()
        compiled = compiler.compile_structural_tag(stag)
        print(f"grammar compile time: {time.monotonic() - t0:.2f}s")

        processor = xgr.contrib.hf.LogitsProcessor(compiled)
        t0 = time.monotonic()
        out = model.generate(
            **inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
        )
        gen = tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False)
        print(f"generate time: {time.monotonic() - t0:.1f}s")
        print(f"--- constrained output ({label}) ---")
        print(repr(gen))
        results[label] = gen

    print("\n=== unconstrained (alignment check) ===")
    out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
    unconstrained = tokenizer.decode(
        out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False
    )
    print(repr(unconstrained))
    results["unconstrained"] = unconstrained

    print("\n=== verification ===")
    ok = True
    required_out = results["tool_choice=required, reasoning=False"]
    # strip trailing special tokens after the tool call for the checks
    body = required_out.split("<tool_call|>")[0] + "<tool_call|>" if "<tool_call|>" in required_out else required_out
    checks = [
        ("required: contains <|tool_call>call:get_weather{", "<|tool_call>call:get_weather{" in body),
        ("required: closes with }<tool_call|>", body.rstrip().endswith("<tool_call|>")),
        ('required: city uses <|"|> delimiters', '<|"|>' in body),
        ("required: no JSON-quoted args", '"city"' not in body),
    ]
    stag_req = xgr.get_model_structural_tag(
        model="gemma_4", tools=TOOLS, tool_choice="required", reasoning=False
    )
    g = xgr.Grammar.from_structural_tag(stag_req)
    checks.append(("required: output accepted by grammar", _is_grammar_accept_string(g, body)))

    for name, passed in checks:
        print(f"  [{'PASS' if passed else 'FAIL'}] {name}")
        ok &= passed

    print("\nRESULT:", "ALL PASS" if ok else "SOME CHECKS FAILED")
    print("\nJSON summary:")
    print(json.dumps({k: v for k, v in results.items()}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()