"""Single-turn chat with TinyLlama-1.1B-Chat INT4 via ONNX Runtime GenAI.

Writes ``predictions.json`` next to this script.
"""

from __future__ import annotations

import json
from pathlib import Path

import onnxruntime_genai as og
from jinja2 import Environment
from tokenizers import Tokenizer


# -- Configuration -------------------------------------------------------------
# HellaSwag-style commonsense continuation prompt. Small chat models
# (~1B) generate cleanest output on well-scoped everyday scenarios (see
# TinyLlama's 59% HellaSwag acc_norm) vs open-ended factual prose which
# tends to hallucinate and loop.
DEFAULT_PROMPT = (
    "A woman is in the kitchen making pancakes. She pours the batter onto "
    "a hot pan and waits for bubbles to appear on the surface. Once the "
    "bubbles pop, she"
)
MAX_LENGTH = 256  # absolute token budget (prompt + decoded)
DO_SAMPLE = False  # greedy decode for reproducibility
TEMPERATURE = 0.0  # ignored unless do_sample=True


def _raise_exception(msg: str) -> None:
    """Bridge for the chat template's ``raise_exception`` helper."""
    raise RuntimeError(msg)


def render_chat_template(
    template_path: Path, bundle_dir: Path, user_prompt: str
) -> str:
    """Render ``chat_template.jinja`` with a single user turn.

    Reads BOS/EOS tokens from ``tokenizer_config.json`` so the rendered string
    matches the actual tokenizer's special tokens. ``add_generation_prompt=True``
    appends the assistant header so the model continues from there.
    """
    template_src = template_path.read_text(encoding="utf-8")
    with (bundle_dir / "tokenizer_config.json").open() as f:
        tok_cfg = json.load(f)

    env = Environment(trim_blocks=True, lstrip_blocks=True, autoescape=False)
    env.globals["raise_exception"] = _raise_exception
    template = env.from_string(template_src)

    return template.render(
        messages=[{"role": "user", "content": user_prompt}],
        bos_token=tok_cfg.get("bos_token", "<s>"),
        eos_token=tok_cfg.get("eos_token", "</s>"),
        add_generation_prompt=True,
    )


def load_model(bundle_dir: Path) -> tuple[og.Model, Tokenizer]:
    """Load the ONNX Runtime GenAI model and the bundled HF tokenizer."""
    model = og.Model(str(bundle_dir))
    tokenizer = Tokenizer.from_file(str(bundle_dir / "tokenizer.json"))
    return model, tokenizer


def encode_prompt(tokenizer: Tokenizer, bundle_dir: Path, prompt: str) -> list[int]:
    """Tokenise so the model sees exactly one BOS token.

    Whether the prompt already carries BOS is a property of the chat template,
    not of the model: Llama-3.x templates emit ``bos_token`` themselves,
    TinyLlama's does not, and a base model has no template at all. Derive the
    ``add_special_tokens`` value from the prompt rather than hardcoding it, then
    assert the invariant so a template or tokenizer change fails loudly instead
    of silently degrading generation quality.
    """
    bos = json.loads((bundle_dir / "tokenizer_config.json").read_text())["bos_token"]
    ids = tokenizer.encode(prompt, add_special_tokens=not prompt.startswith(bos)).ids
    bos_id = tokenizer.token_to_id(bos)
    if not ids or ids[0] != bos_id or ids.count(bos_id) != 1:
        raise RuntimeError(
            f"prompt must carry exactly one leading {bos!r}; got {ids.count(bos_id)}"
        )
    return ids


def generate(
    model: og.Model,
    tokenizer: Tokenizer,
    bundle_dir: Path,
    chat_prompt: str,
) -> tuple[str, int, int]:
    """Run greedy decode for one chat turn.

    Returns ``(decoded_response, prompt_token_count, generated_token_count)``.
    """
    input_ids = encode_prompt(tokenizer, bundle_dir, chat_prompt)

    params = og.GeneratorParams(model)
    params.set_search_options(
        max_length=MAX_LENGTH,
        do_sample=DO_SAMPLE,
        temperature=TEMPERATURE,
    )

    generator = og.Generator(model, params)
    generator.append_tokens(input_ids)
    while not generator.is_done():
        generator.generate_next_token()

    full_ids = list(generator.get_sequence(0))
    response_ids = full_ids[len(input_ids) :]
    decoded = tokenizer.decode(response_ids, skip_special_tokens=True)
    return decoded, len(input_ids), len(response_ids)


def save_results(
    bundle_dir: Path,
    user_prompt: str,
    response: str,
    prompt_tokens: int,
) -> Path:
    """Persist the prompt/response pair as ``predictions.json``."""
    output_path = bundle_dir / "predictions.json"
    payload = {
        "task": "TinyLlama-1.1B-Chat — single-turn chat",
        "chat_template": "chat_template.jinja",
        "results": [
            {
                "user": user_prompt,
                "assistant": response,
                "prompt_tokens": prompt_tokens,
            }
        ],
        "max_length": MAX_LENGTH,
        "do_sample": DO_SAMPLE,
        "temperature": TEMPERATURE,
    }
    output_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False))
    return output_path


def main() -> None:
    bundle_dir = Path(__file__).resolve().parent
    template_path = bundle_dir / "chat_template.jinja"

    print(f"Loading model from: {bundle_dir}")
    model, tokenizer = load_model(bundle_dir)

    chat_prompt = render_chat_template(template_path, bundle_dir, DEFAULT_PROMPT)
    print(f"\nPrompt: {DEFAULT_PROMPT}\n")
    print("Generating...")
    response, prompt_tokens, generated_tokens = generate(model, tokenizer, bundle_dir, chat_prompt)

    print("\n--- Response ---")
    print(response)
    print("--- /Response ---")
    print(
        f"\nPrompt tokens: {prompt_tokens} | " f"Generated tokens: {generated_tokens}"
    )

    saved = save_results(bundle_dir, DEFAULT_PROMPT, response, prompt_tokens)
    print(f"Saved: {saved}")


if __name__ == "__main__":
    main()