"""Chat / text generation with the ONNX KnowLine model (raw onnxruntime, no torch).

    python generate.py --model ./model-cuda-int8 --cuda --prompt "Hello"   # CUDA, int8
    python generate.py --model ./model-cuda-int4 --cuda --prompt "Hello"   # CUDA, int4
    python generate.py --model ./model-cpu-int8 --prompt "Hello"   # CPU, int8
    python generate.py --model ./model-cpu --interactive
    python generate.py --model ./model-cpu --prompt "Hi" --max-new-tokens 100 --no-chat-template

Uses onnx_raw.py (plain onnxruntime prefill + greedy decode). onnxruntime-genai
0.19.0 cannot load this Qwen3.5 export on CPU, so it is intentionally not used.
"""

import argparse
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

from onnx_raw import OnnxLM  # noqa: E402


def eos_ids(model_dir):
    import json

    try:
        cfg = json.loads((Path(model_dir) / "genai_config.json").read_text())
        eos = cfg["model"]["eos_token_id"]
        return tuple(eos if isinstance(eos, list) else [eos])
    except (OSError, ValueError, KeyError):
        return (248046, 248044)  # <|im_end|>, <|endoftext|> (Qwen3.5)


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", default=str(HERE / "model-cpu"), help="build dir (model.onnx lives inside)")
    p.add_argument("--onnx", default="model.onnx", help="onnx file name inside --model")
    p.add_argument("--cuda", action="store_true", help="use CUDAExecutionProvider")
    p.add_argument("--prompt", default=None)
    p.add_argument("--interactive", action="store_true")
    p.add_argument("--max-new-tokens", type=int, default=512)
    p.add_argument("--no-chat-template", action="store_true", help="use the prompt verbatim (no chat template)")
    p.add_argument("--threads", type=int, default=0, help="ORT intra-op threads (0 = default)")
    a = p.parse_args()

    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(str(HERE / "orig_tokenizer"), local_files_only=True)
    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if a.cuda else ["CPUExecutionProvider"]
    print(f"Loading {a.model}/{a.onnx} {providers[0]} ...", flush=True)
    lm = OnnxLM(str(Path(a.model) / a.onnx), providers=providers, intra_threads=a.threads)
    eos = eos_ids(a.model)

    def run(prompt):
        if a.no_chat_template:
            ids = tok.encode(prompt, add_special_tokens=False)
        else:
            text = tok.apply_chat_template([{"role": "user", "content": prompt}],
                                           tokenize=False, add_generation_prompt=True, enable_thinking=False)
            ids = tok.encode(text, add_special_tokens=False)
        print(f"  prompt tokens: {len(ids)}", flush=True)
        t0 = time.perf_counter()
        n_first = {"done": False}
        t_first = {}

        def on_token(t):
            if not n_first["done"]:
                n_first["done"] = True
                t_first["t"] = time.perf_counter() - t0
            sys.stdout.write(tok.decode([t]))
            sys.stdout.flush()

        out = lm.generate(ids, a.max_new_tokens, eos, on_token=on_token)
        total = time.perf_counter() - t0
        print()
        n = len(out)
        print(f"  tokens={n} TTFT={t_first.get('t', total) * 1000:.0f}ms "
              f"TPS={max(n - 1, 1) / max(total - t_first.get('t', 0), 1e-9):.1f} total={total:.1f}s")

    if a.interactive:
        print("Interactive mode ('quit' to stop).")
        while True:
            try:
                prompt = input("You: ").strip()
            except EOFError:
                break
            if prompt.lower() in ("quit", "exit") or not prompt:
                break
            run(prompt)
            print("-" * 50)
    elif a.prompt:
        run(a.prompt)
    else:
        p.print_help()


if __name__ == "__main__":
    main()