"""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 os
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"
QUANTIZE_INT8 = os.environ.get("QUANTIZE_INT8", "0") == "1"

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)
    if QUANTIZE_INT8:
        from torchao.quantization import Int8WeightOnlyConfig
        from transformers import TorchAoConfig

        print("=== loading with torchao int8 weight-only quantization ===")
        quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=torch.bfloat16,
            device_map="mps",
            quantization_config=quantization_config,
        )
    else:
        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}")

    # reasoning=False prompt/inputs; also used below for the unconstrained alignment
    # check, which is compared against the tool_choice=required, reasoning=False case.
    prompt = tokenizer.apply_chat_template(
        MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False, enable_thinking=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")

        # Rebuild the prompt so <|think|> is present exactly when this iteration's
        # grammar expects a thought block (enable_thinking must match `reasoning`).
        iter_prompt = tokenizer.apply_chat_template(
            MESSAGES,
            tools=TOOLS,
            add_generation_prompt=True,
            tokenize=False,
            enable_thinking=reasoning,
        )
        iter_inputs = tokenizer(iter_prompt, return_tensors="pt", add_special_tokens=False).to(
            "mps"
        )

        processor = xgr.contrib.hf.LogitsProcessor(compiled)
        t0 = time.monotonic()
        out = model.generate(
            **iter_inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
        )
        gen = tokenizer.decode(
            out[0][iter_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()