#!/usr/bin/env python
"""
End-to-end MCQ evaluation with LLaVA-NeXT-Video-7B-DPO.

Usage:
    python run.py --frame-budgets 0 64 128 256
"""
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 numpy as np
import torch
from transformers import LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor

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_letters
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_letters
from util.video import resolve_video_path

from decord import VideoReader, cpu


MODEL_ID = DEFAULT_MODELS['llava_next_video']
EXPERIMENT_NAME = 'llava_next_video'
LLAVA_NEXT_MAX_FRAMES = 64


def _sample_frames(video_path: Path, n: int) -> np.ndarray | None:
    """Return (n, H, W, 3) uint8 array of uniformly sampled frames, or None for n=0."""
    if n == 0:
        return None
    vr = VideoReader(str(video_path), ctx=cpu(0))
    idx = np.linspace(0, len(vr) - 1, n, dtype=int)
    return vr.get_batch(idx).asnumpy()


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

    over = [n for n in args.frame_budgets if n > LLAVA_NEXT_MAX_FRAMES]
    if over:
        sys.exit(
            f'ERROR: LLaVA-Next-Video has max_position_embeddings ~ 10250; '
            f'frame budgets {over} would all overflow.\n'
            f'  Cap is {LLAVA_NEXT_MAX_FRAMES} frames (each ~145 visual tokens).\n'
        )

    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}...')
    processor = LlavaNextVideoProcessor.from_pretrained(MODEL_ID)
    model = LlavaNextVideoForConditionalGeneration.from_pretrained(
        MODEL_ID, torch_dtype=torch.bfloat16,
        device_map=args.device_map, attn_implementation='flash_attention_2',
        low_cpu_mem_usage=True,
    )
    model.eval()
    print(f'Loaded on {model.device}, dtype={model.dtype}')

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


def _run_one(df, num_frames: int, args, out_dir: Path, processor, model) -> 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 {'video': None}   # text-only, no video
        try:
            vpath = resolve_video_path(vid)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e}')
            return {}
        try:
            video = _sample_frames(vpath, num_frames)
        except Exception as e:
            print(f'  Frame sample error {vid}: {type(e).__name__}: {e!r}')
            return {}
        return {'video': video}

    def teardown_movie_ctx(ctx):
        ctx.pop('video', None)
        import gc; gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def call_fn(row, ctx):
        # Text-only vs video branches
        video = ctx.get('video')
        if num_frames > 0 and video is None:
            return ''

        header = HEADER_TEXT_ONLY if num_frames == 0 else HEADER_VIDEO_FRAMES.format(n_frames=num_frames)
        prompt_text = mcq_prompt_letters(row, header=header)

        if video is None:
            conversation = [{'role': 'user', 'content': [{'type': 'text', 'text': prompt_text}]}]
            prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
            inputs = processor(text=prompt, return_tensors='pt')
        else:
            conversation = [{'role': 'user', 'content': [
                {'type': 'video'}, {'type': 'text', 'text': prompt_text},
            ]}]
            prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
            inputs = processor(text=prompt, videos=[video], return_tensors='pt')

        inputs = {k: (v.to(model.device, dtype=model.dtype) if v.dtype.is_floating_point else v.to(model.device))
                  if hasattr(v, 'to') else v
                  for k, v in inputs.items()}

        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
        raw = processor.decode(out[0, inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        return parse_mcq_letters(raw)

    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        teardown_movie_ctx=teardown_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':   'llava_next_video',
            'mode':       'text-only' if num_frames == 0 else 'video',
        },
        checkpoint_every=args.checkpoint_every,
        # In text-only mode, build_movie_ctx returns {} for skip_if_no_ctx check —
        # but since text-only doesn't need video, we pass `skip_if_no_ctx=False` to allow it.
        skip_if_no_ctx=(num_frames > 0),
    )


if __name__ == '__main__':
    main()