#!/usr/bin/env python
"""
End-to-end MCQ evaluation with Qwen3-VL-8B-Instruct.

Usage:
    python run.py --frame-budgets 0 64 128 256
    python run.py --frame-budgets 128 --checkpoint-every 25
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

# Make util.* importable without a package install.
_PROJECT = Path(__file__).resolve().parents[3]  # chnb/
sys.path.insert(0, str(_PROJECT))

from util.backbones import init_backbone, call_backbone
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 HEADER_TEXT_ONLY, HEADER_VIDEO_FRAMES, mcq_prompt_json
from util.video import load_uniform_frames


EXPERIMENT_NAME = 'qwen3vl'
MODEL_ID = DEFAULT_MODELS['qwen3vl']


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.strip().split('\n\n', 1)[0])
    ap.add_argument('--frame-budgets', type=int, nargs='+', default=[0, 64, 128, 256])
    ap.add_argument('--benchmark', type=str, default=str(BENCHMARK_PATH))
    ap.add_argument('--max-new-tokens', type=int, default=1024)
    ap.add_argument('--temperature', type=float, default=0.1)
    ap.add_argument('--checkpoint-every', type=int, default=50)
    ap.add_argument('--device-map', type=str, default='auto')
    ap.add_argument('--min-pixels', type=int, default=256 * 28 * 28)
    ap.add_argument('--max-pixels', type=int, default=512 * 28 * 28)
    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}...')
    init_backbone('qwen3vl', device_map=args.device_map,
                  min_pixels=args.min_pixels, max_pixels=args.max_pixels)
    print('Ready.')

    for num_frames in args.frame_budgets:
        _run_one_budget(df, num_frames, args, out_dir)


def _run_one_budget(df, num_frames: int, args, out_dir: Path) -> None:
    mode_label = 'text-only' if num_frames == 0 else f'video ({num_frames} frames)'
    print(f'\n=== {EXPERIMENT_NAME} f={num_frames} [{mode_label}] ===')

    def build_movie_ctx(vid, movie_df):
        if num_frames == 0:
            return {'frames': None}   # text-only, no video
        try:
            frames = load_uniform_frames(vid, num_frames, strict=True)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e}')
            return {}
        return {'frames': frames}

    def teardown_movie_ctx(ctx):
        ctx.pop('frames', None)
        import gc; gc.collect()
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except ImportError:
            pass

    def call_fn(row, ctx):
        frames = ctx.get('frames')
        if num_frames > 0 and not frames:
            return ''
        header = HEADER_TEXT_ONLY if num_frames == 0 else HEADER_VIDEO_FRAMES.format(n_frames=num_frames)
        prompt = mcq_prompt_json(row, header=header)
        raw = call_backbone(
            'qwen3vl', prompt, images=frames,
            system_instruction='You are a helpful assistant answering questions about a movie. Output valid JSON.',
            json_mode=True,
            max_output_tokens=args.max_new_tokens,
            temperature=args.temperature,
        )
        return parse_mcq_json(raw, qid=row['qid'])

    results_path = out_dir / results_filename_for_budget(num_frames, MODEL_ID)
    run_sweep(
        df,
        call_fn=call_fn,
        results_path=results_path,
        experiment_name=f'{EXPERIMENT_NAME} f={num_frames}',
        build_movie_ctx=build_movie_ctx,
        teardown_movie_ctx=teardown_movie_ctx,
        extra_row_fields={
            'num_frames': num_frames,
            'model':      MODEL_ID,
            'backbone':   'qwen3vl',
            'mode':       'text-only' if num_frames == 0 else 'video',
        },
        checkpoint_every=args.checkpoint_every,
        # Text-only returns {'frames': None} (truthy dict), so skip_if_no_ctx only
        # fires for video budgets where frame loading actually failed.
        skip_if_no_ctx=(num_frames > 0),
    )


if __name__ == '__main__':
    main()