#!/usr/bin/env python3
"""Create one exact-length, prefix-disjoint prefill request."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--input-tokens", type=int, default=8192)
    parser.add_argument("--output-tokens", type=int, default=2)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if min(args.input_tokens, args.output_tokens) <= 0:
        raise ValueError("token counts must be positive")
    if args.input_tokens + args.output_tokens > 40960:
        raise ValueError("request exceeds the server max model length")

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    special = set(tokenizer.all_special_ids)
    candidates = [
        token for token in range(tokenizer.vocab_size) if token not in special
    ]
    if len(candidates) < 2:
        raise ValueError("tokenizer has too few non-special token IDs")

    body = {
        "model": "qwen30-prefill-profile",
        "prompt": [candidates[1], *([candidates[0]] * (args.input_tokens - 1))],
        "min_tokens": args.output_tokens,
        "max_tokens": args.output_tokens,
        "ignore_eos": True,
        "temperature": 0,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(body, separators=(",", ":")) + "\n")


if __name__ == "__main__":
    main()