"""Download/run the final checkpoint, printing one JSON snapshot per readout."""

import argparse
import json
import os
from pathlib import Path

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
import torch
from huggingface_hub import snapshot_download

from scripts.demo.inference import DemoRequest, stream_inference
from scripts.recurrent_qwen.checkpoint import load_recurrent_checkpoint


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--checkpoint", type=Path)
    source.add_argument("--model", help="Hugging Face model repository ID")
    parser.add_argument("--revision", default="main")
    parser.add_argument("--request", required=True, type=Path, help="JSON exported from the demo")
    parser.add_argument("--device", choices=("cpu", "cuda", "mps"), default="cpu")
    parser.add_argument("--download", action="store_true", help="Allow explicit Hub/base downloads")
    parser.add_argument("--max-steps", type=int, default=256)
    parser.add_argument("--loop-cap", type=int, default=272)
    args = parser.parse_args()
    request = DemoRequest.parse(json.loads(args.request.read_text()), max_steps=args.max_steps)
    if args.loop_cap < request.steps:
        parser.error("--loop-cap must cover the requested steps")
    path = args.checkpoint or Path(snapshot_download(args.model, revision=args.revision,
                                                   local_files_only=not args.download))
    torch.set_num_threads(4)
    torch.use_deterministic_algorithms(True)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    model, tokenizer, spec = load_recurrent_checkpoint(path, device=args.device, download=args.download)
    model.config._attn_implementation = "sdpa"
    model.requires_grad_(False)
    for frame in stream_inference(model, tokenizer, spec, request,
                                   checkpoint_id=args.model or str(path), loop_cap=args.loop_cap):
        print(json.dumps(frame, allow_nan=False), flush=True)


if __name__ == "__main__":
    main()