#!/usr/bin/env python
"""
End-to-end MCQ evaluation with Gemini 3 Flash over pre-extracted uniform frames.

Calls Vertex Gemini via `util.backbones.call_backbone('gemini', ...)`.

Usage:
    python run.py --frame-budgets 64 128 256
    python run.py --frame-budgets 128 --model gemini-3-flash-preview
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

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

from util.backbones import call_backbone, init_backbone
from util.benchmark import load_benchmark
from util.config import (
    DEFAULT_MODELS,
    DEFAULT_VERTEX_LOCATION,
    GOOGLE_CLOUD_PROJECT_ENV,
)
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 = 'gemini3flash'


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--frame-budgets', type=int, nargs='+', default=[0, 128, 256])
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--model', default=DEFAULT_MODELS['gemini'])
    ap.add_argument('--project', default=os.environ.get(GOOGLE_CLOUD_PROJECT_ENV))
    ap.add_argument('--location', default=DEFAULT_VERTEX_LOCATION)
    ap.add_argument('--max-output-tokens', type=int, default=24576)
    ap.add_argument('--temperature', type=float, default=0.0)
    ap.add_argument('--checkpoint-every', type=int, default=50)
    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')

    init_backbone('gemini', project_id=args.project, location=args.location, model_id=args.model)
    print(f'Gemini ready: {args.model}')

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


def _run_one(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 frames needed
        try:
            frames = load_uniform_frames(vid, num_frames, strict=True)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e}')
            return {}
        return {'frames': frames}

    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(
            'gemini', prompt, images=frames,
            system_instruction='You answer movie MCQs. Output only valid JSON.',
            json_mode=True,
            max_output_tokens=args.max_output_tokens,
            temperature=args.temperature,
        )
        return parse_mcq_json(raw, qid=row['qid'])

    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        results_path=out_dir / results_filename_for_budget(num_frames, args.model),
        experiment_name=f'{EXPERIMENT_NAME} f={num_frames}',
        extra_row_fields={
            'num_frames': num_frames,
            'model':      args.model,
            'backbone':   'gemini',
            'mode':       'text-only' if num_frames == 0 else 'video',
        },
        checkpoint_every=args.checkpoint_every,
        skip_if_no_ctx=(num_frames > 0),
    )


if __name__ == '__main__':
    main()