#!/usr/bin/env python3
"""End-to-end inference with DCVC-RT-selected patches on the release model.

Two modes:

  * ``--asset_dir DIR``  use canvases already written by ``precompute_dcvc_rt.py``.
  * ``--video PATH``     precompute on the fly (calls ``precompute_dcvc_rt.py``)
                         into ``--work_dir`` then runs inference.

Example::

    python neural_codec/infer_dcvc_rt.py \
        --video clip.mp4 \
        --question "Describe what happens in this video." \
        --max_pixels 150000 --qp 21
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path

import torch

_HERE = os.path.dirname(os.path.abspath(__file__))
_MODEL_DEFAULT = os.path.dirname(_HERE)  # Mage-VL-Ported/ (this file lives in neural_codec/)

sys.path.insert(0, _HERE)
from codec_loader import build_inputs_from_assets  # noqa: E402


def _precompute_inline(args) -> str:
    key = Path(args.video).stem
    out_root = args.work_dir
    asset_dir = os.path.join(out_root, "assets", key)
    py = args.precompute_python or sys.executable
    cmd = [
        py, os.path.join(_HERE, "precompute_dcvc_rt.py"),
        "--videos", args.video, "--out_root", out_root,
        "--intra_tar", args.intra_tar, "--inter_tar", args.inter_tar,
        "--target_canvas", str(args.target_canvas), "--seq_len_frames", str(args.seq_len_frames),
        "--max_pixels", str(args.max_pixels), "--qp", str(args.qp),
        "--dcvc_max_side", str(args.dcvc_max_side), "--cuda_idx", str(args.gpu),
        "--num_workers", "1",
    ]
    if args.canvas_token_side:
        cmd += ["--canvas_token_side", str(args.canvas_token_side)]
    print("[infer] precomputing:", " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)
    return asset_dir


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", default=_MODEL_DEFAULT)
    src = ap.add_mutually_exclusive_group(required=True)
    src.add_argument("--asset_dir", help="precomputed assets/<key> dir")
    src.add_argument("--video", help="video to precompute + run")
    ap.add_argument("--question", default="Describe this video in detail.")
    ap.add_argument("--max_new_tokens", type=int, default=256)
    ap.add_argument("--max_pixels", type=int, default=150000)
    ap.add_argument("--gpu", type=int, default=0)
    # inline-precompute knobs (only for --video)
    ap.add_argument("--work_dir", default="/tmp/dcvc_rt_infer")
    ap.add_argument("--precompute_python", default=None,
                    help="interpreter for the DCVC-RT precompute subprocess (needs the "
                         "'codec' env); default: the current interpreter (sys.executable)")
    ap.add_argument("--intra_tar", default=os.path.join(_HERE, "dcvc_rt_intra.tar"))
    ap.add_argument("--inter_tar", default=os.path.join(_HERE, "dcvc_rt_inter.tar"))
    ap.add_argument("--target_canvas", type=int, default=32)
    ap.add_argument("--seq_len_frames", type=int, default=64)
    ap.add_argument("--canvas_token_side", type=int, default=None)
    ap.add_argument("--qp", type=int, default=21)
    ap.add_argument("--dcvc_max_side", type=int, default=512)
    args = ap.parse_args()

    asset_dir = args.asset_dir or _precompute_inline(args)

    sys.path.insert(0, os.path.join(args.model, "processor"))   # MageVLProcessor + siblings
    from processing_magevl import MageVLProcessor
    from transformers import AutoModelForCausalLM

    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
    print("[infer] loading model:", args.model, flush=True)
    processor = MageVLProcessor.from_pretrained(args.model)       # subdir-aware
    model = AutoModelForCausalLM.from_pretrained(
        os.path.join(args.model, "transformer"), trust_remote_code=True, torch_dtype="auto"
    ).to(device).eval()

    # Build the chat-templated text with a video span for rewrite to replace.
    messages = [{"role": "user", "content": [
        {"type": "video"},
        {"type": "text", "text": args.question},
    ]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = build_inputs_from_assets(processor, asset_dir, text,
                                      max_pixels=args.max_pixels, device=device)
    # match vision dtype to model
    inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)

    print("[infer] input_ids", tuple(inputs["input_ids"].shape),
          "pixel_values", tuple(inputs["pixel_values"].shape),
          "grid_thw", inputs["image_grid_thw"].tolist(),
          "patch_positions", tuple(inputs["patch_positions"].shape), flush=True)

    with torch.inference_mode():
        gen = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
    new_tokens = gen[0, inputs["input_ids"].shape[1]:]
    answer = processor.tokenizer.decode(new_tokens, skip_special_tokens=True)
    print("\n================ ANSWER ================\n" + answer.strip() + "\n")


if __name__ == "__main__":
    main()