#!/usr/bin/env python
"""
End-to-end MCQ evaluation with VideoChat-Flash-Qwen2_5-7B-1M_res224.

Usage:
    python run.py --frame-budgets 128 256
    CUDA_VISIBLE_DEVICES=0,1,2,3 python run.py --frame-budgets 1024
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

_PROJECT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(_PROJECT))

import torch
from transformers import AutoModel, AutoTokenizer

from util.benchmark import load_benchmark
from util.config import DEFAULT_MODELS
from util.eval_loop import run_sweep
from util.parsers import parse_mcq_json
from util.paths import (
    BENCHMARK_PATH, RESULTS_DIR, ensure_dir,
    results_filename_for_budget,
)
from util.prompts import mcq_prompt_json
from util.video import resolve_video_path


MODEL_ID = DEFAULT_MODELS['videochat']
EXPERIMENT_NAME = 'videochat'


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--frame-budgets', type=int, nargs='+', default=[64, 128, 256])
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--max-new-tokens', type=int, default=4096)
    ap.add_argument('--checkpoint-every', type=int, default=50)
    ap.add_argument('--device-map', default='auto')
    args = ap.parse_args()

    out_dir = ensure_dir(RESULTS_DIR / EXPERIMENT_NAME)
    df = load_benchmark(args.benchmark)
    print(f'Loaded {len(df)} questions across {df["video_file_id"].nunique()} movies')

    print(f'Loading {MODEL_ID}...')
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    # VideoChat's custom code wants bfloat16 + cuda; device_map='auto' lets it
    # shard across visible GPUs for larger frame budgets (1024+).
    model = AutoModel.from_pretrained(
        MODEL_ID, trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        device_map=args.device_map,
    ).eval()
    print(f'Loaded.')

    for num_frames in args.frame_budgets:
        if num_frames == 0:
            print(f'  skipping f=0: VideoChat has no clean text-only path. '
                  f'Use qwen3vl --frame-budgets 0 for the LLM-only baseline.')
            continue
        _run_one(df, num_frames, args, out_dir, tokenizer, model)


def _run_one(df, num_frames: int, args, out_dir: Path, tokenizer, model) -> None:
    print(f'\n=== {EXPERIMENT_NAME} f={num_frames} ===')

    def build_movie_ctx(vid, movie_df):
        try:
            vpath = resolve_video_path(vid)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e}')
            return {}
        return {'video_path': str(vpath)}

    # Track raw output of the first few unparseable responses so the user
    # can see what the model actually said when parse_mcq_json returns ''.
    # The eval_loop only logs the empty pred event; it doesn't carry the
    # raw text. We print it here right before returning '' so it shows up
    # immediately above the eval-loop's `[empty pred]` line.
    _raw_logged = {'count': 0, 'limit': 8}

    def call_fn(row, ctx):
        vpath = ctx.get('video_path')
        if not vpath:
            return ''
        # VideoChat handles frame sampling internally — no images passed.
        prompt = mcq_prompt_json(row)
        try:
            output = model.chat(
                video_path=vpath, tokenizer=tokenizer,
                user_prompt=prompt, return_history=False,
                max_num_frames=num_frames,
                generation_config=dict(
                    do_sample=False, temperature=0.0,
                    max_new_tokens=args.max_new_tokens, num_beams=1,
                ),
            )
        except Exception as e:
            # VideoChat is fussy with high frame counts; let the empty pred retry
            print(f'  chat error {row["qid"]}: {type(e).__name__}: {e!r}')
            return ''
        raw = output if isinstance(output, str) else (output[0] if output else '')
        pred = parse_mcq_json(raw, qid=row['qid'])
        if not pred and _raw_logged['count'] < _raw_logged['limit']:
            preview = (raw[:300] if isinstance(raw, str) else repr(raw)[:300])
            preview = preview.replace('\n', ' ⏎ ')
            print(f'  [parse fail] qid={row["qid"]} raw={preview!r}')
            _raw_logged['count'] += 1
        return pred

    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        results_path=out_dir / results_filename_for_budget(num_frames, MODEL_ID),
        experiment_name=f'{EXPERIMENT_NAME} f={num_frames}',
        extra_row_fields={
            'num_frames': num_frames,
            'model':      MODEL_ID,
            'backbone':   'videochat',
        },
        checkpoint_every=args.checkpoint_every,
    )


if __name__ == '__main__':
    main()