"""Génération locale avec le checkpoint PyTorch Ys-17M."""
from __future__ import annotations

import argparse
import hashlib
from pathlib import Path

import torch
from tokenizers import Tokenizer

from model import GPT, ModelConfig


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", default="Il était une fois")
    parser.add_argument("--max-new-tokens", type=int, default=80)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    args = parser.parse_args()

    selected = args.device
    if selected == "auto":
        selected = "cuda" if torch.cuda.is_available() else "cpu"
    if selected == "cuda" and not torch.cuda.is_available():
        parser.error("CUDA indisponible : utilise --device cpu.")
    device = torch.device(selected)
    torch.set_num_threads(min(8, torch.get_num_threads()))
    torch.manual_seed(args.seed)

    directory = Path(__file__).resolve().parent
    checkpoint = torch.load(directory / "best.pt", map_location="cpu", weights_only=True)
    tokenizer_path = directory / "tokenizer.json"
    actual_hash = hashlib.sha256(tokenizer_path.read_bytes()).hexdigest()
    if actual_hash != checkpoint["tokenizer_hash"]:
        raise ValueError("Le tokenizer ne correspond pas au checkpoint Ys-17M.")
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    model = GPT(ModelConfig(**checkpoint["model_config"]))
    model.load_state_dict(checkpoint["model"], strict=True)
    model.to(device).eval()

    ids = tokenizer.encode(args.prompt).ids or [tokenizer.token_to_id("<|bos|>")]
    inputs = torch.tensor([ids], dtype=torch.long, device=device)
    result = model.generate(
        inputs,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        eos_id=tokenizer.token_to_id("<|eos|>"),
    )
    print(tokenizer.decode(result[0].tolist(), skip_special_tokens=True))


if __name__ == "__main__":
    main()