#!/usr/bin/env python3
"""
generate.py — CLI for the inference.py library bundled in this repo.

One-shot:
    python generate.py --prompt "<s><|yor_Latn|>Kaabo..." -n 4

Interactive REPL:
    python generate.py --interactive

Prompt format: prefix with <s> (BOS) and a language token like <|yor_Latn|>.
Available language tokens are listed in tokenizer.json.
"""

import argparse
import sys
from pathlib import Path

from inference import load_model, generate


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--repo-dir", default=".", help="Path to the repo (default: cwd).")
    ap.add_argument("--prompt", default=None, help="Single prompt to complete.")
    ap.add_argument("--interactive", action="store_true",
                    help="REPL mode. Type prompts; empty line or Ctrl-D exits.")
    ap.add_argument("-n", "--num-sequences", type=int, default=1)
    ap.add_argument("--max-new-tokens", type=int, default=128)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top-p", type=float, default=0.9)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--device", default=None,
                    help="cuda / cpu (default: cuda if available).")
    args = ap.parse_args()

    if not args.prompt and not args.interactive:
        print("[error] specify --prompt or --interactive", file=sys.stderr); sys.exit(1)

    print(f"[load] {args.repo_dir}")
    model, tokenizer, config = load_model(args.repo_dir, device=args.device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[load] {n_params/1e6:.1f}M params, dtype={next(model.parameters()).dtype}, "
          f"vocab={config.vocab_size}, T={config.max_seq_len}")

    def run_one(prompt: str):
        out = generate(
            model, tokenizer, prompt,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature, top_p=args.top_p,
            num_sequences=args.num_sequences, seed=args.seed,
        )
        for i, c in enumerate(out):
            print(f"\n[sample {i}]")
            print(c)

    if args.prompt:
        run_one(args.prompt)
        return

    print("\nInteractive mode. Empty line or Ctrl-D to exit.")
    print('Prompt format: <s><|lang_Script|>your text here')
    while True:
        try:
            p = input("\n>>> ").rstrip()
        except (EOFError, KeyboardInterrupt):
            print(); break
        if not p:
            break
        run_one(p)


if __name__ == "__main__":
    main()