#!/usr/bin/env python
"""
VideoAgent-style agentic frame retrieval for MCQ evaluation.

Cf. https://github.com/YueFan1014/VideoAgent

Usage:
    python run.py --backbone gemini
    python run.py --backbone gpt --initial-frames 8 --max-iterations 3
    python run.py --backbone qwen3vl --pool-size 64
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

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

import numpy as np
from PIL import Image

from util.backbones import call_backbone, init_backbone
from util.benchmark import load_benchmark
from util.eval_loop import run_sweep
from util.frames import get_clip_embeddings, retrieve_frames_by_text
from util.parsers import norm_qid, parse_json_flexible, parse_mcq_json
from util.paths import (
    BENCHMARK_PATH, RESULTS_DIR, ensure_dir, get_clip_path,
    results_filename,
)
from util.video import load_uniform_frames


EXPERIMENT_NAME = 'videoagent'
SUPPORTED_BACKBONES = ('gemini', 'gpt', 'claude', 'qwen3vl', 'glm4v')

# Pool sizes match the end-to-end frame-budget axis. The backbone only sees
# `initial_frames + retrievals` (~20) images per call, so the pool size
# itself doesn't trigger OOM / context / 100-image limits.
DEFAULT_POOL_SIZES = [64, 128, 256]


def _init_backbone_guard(backbone: str) -> None:
    from util.secrets import require
    if backbone == 'gpt':
        require('OPENAI_API_KEY')
    if backbone == 'claude':
        require('ANTHROPIC_API_KEY')
    if backbone == 'gemini':
        init_backbone('gemini', project_id=os.environ.get('GOOGLE_CLOUD_PROJECT'))
    else:
        init_backbone(backbone)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--backbone', default='gemini', choices=SUPPORTED_BACKBONES,
                    help='Handles both captioning and QA.')
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--pool-sizes', type=int, nargs='+', default=DEFAULT_POOL_SIZES,
                    help='Uniform-frame pool sizes to sweep. Each size is its own '
                         'CSV (results_{backbone}_pool{N}.csv). Match end-to-end frame budgets.')
    ap.add_argument('--initial-frames', type=int, default=8,
                    help='Number of uniformly-sampled frames to start each loop with.')
    ap.add_argument('--max-iterations', type=int, default=3)
    ap.add_argument('--confidence-threshold', type=int, default=3, choices=[1, 2, 3],
                    help='Stop loop when self-eval confidence ≥ threshold (1=insufficient, 3=sufficient).')
    ap.add_argument('--frames-per-retrieval', type=int, default=4)
    ap.add_argument('--max-output-tokens', type=int, default=24576)
    ap.add_argument('--temperature', type=float, default=0.1)
    ap.add_argument('--checkpoint-every', type=int, default=25)
    ap.add_argument('--clip-device', default='cpu',
                    choices=['cpu', 'cuda', 'auto'],
                    help='Where to run CLIP. Default cpu.')
    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'Initializing backbone: {args.backbone}')
    _init_backbone_guard(args.backbone)

    # CLIP loaded once and reused; default CPU because the runtime workload
    # is short text-encoding, not on-the-fly image encoding (preprocess
    # caches the frame embeddings).
    from transformers import CLIPModel, CLIPProcessor
    import torch
    if args.clip_device == 'cuda' and not torch.cuda.is_available():
        sys.exit('ERROR: --clip-device cuda but no CUDA device visible. '
                 'Either drop the flag (defaults to cpu) or check CUDA_VISIBLE_DEVICES.')
    target = args.clip_device
    if target == 'auto':
        target = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'Loading CLIP (openai/clip-vit-large-patch14) on {target}...')
    clip_model = CLIPModel.from_pretrained('openai/clip-vit-large-patch14').eval()
    clip_processor = CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14')
    if target == 'cuda':
        clip_model = clip_model.cuda()
    print(f'CLIP ready on {next(clip_model.parameters()).device}')

    for pool_size in args.pool_sizes:
        args.pool_size = pool_size
        print(f'\n=== pool_size = {pool_size} ===')
        _run_one(df, args, out_dir, clip_model, clip_processor)


def _run_one(df, args, out_dir: Path, clip_model, clip_processor) -> None:
    backbone = args.backbone

    def build_movie_ctx(vid, movie_df):
        try:
            all_frames = load_uniform_frames(vid, args.pool_size, strict=True)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e} — run preprocess --uniform-n {args.pool_size}')
            return {}

        # Pre-computed CLIP embeddings from preprocess; fall back to on-the-fly if missing.
        emb_path = get_clip_path(vid, 'uniform', args.pool_size)
        if emb_path.exists():
            embs = np.load(emb_path)
        else:
            print(f'  computing CLIP embs for {vid} (cached for next run)...')
            embs = get_clip_embeddings(
                all_frames, cache_path=str(emb_path),
                model=clip_model, processor=clip_processor,
            )

        # Synthetic 1-min-per-frame timestamps (pool is uniform samples from full movie).
        # Only used for display in frame labels; not load-bearing.
        n = len(all_frames)
        # Reasonable default: spread timestamps across a ~90-min movie
        timestamps = np.linspace(0, 90 * 60, n)

        return {
            'frames':     all_frames,
            'embeddings': embs,
            'timestamps': timestamps,
        }

    def teardown_movie_ctx(ctx):
        ctx.pop('frames', None)
        ctx.pop('embeddings', 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):
        if not ctx:
            return ''

        final_preds, n_iter, n_frames_used = _agentic_loop(
            question_row=row,
            movie_title=row['movie_title'],
            all_frames=ctx['frames'],
            all_timestamps=ctx['timestamps'],
            clip_embs=ctx['embeddings'],
            args=args,
            clip_model=clip_model,
            clip_processor=clip_processor,
        )
        # Stash per-call telemetry so extra_row_fields can pick them up below.
        row['_n_iter'] = n_iter
        row['_n_frames_used'] = n_frames_used
        return final_preds.get(norm_qid(row['qid']), '')

    model_id = _model_id_for(backbone)
    results_path = out_dir / results_filename(
        backbone=backbone, model_id=model_id,
        config=f'pool_{args.pool_size}',
    )
    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        teardown_movie_ctx=teardown_movie_ctx,
        results_path=results_path,
        experiment_name=f'{EXPERIMENT_NAME} backbone={backbone} pool={args.pool_size}',
        extra_row_fields={
            'backbone':      backbone,
            'model':         _model_id_for(backbone),
            'pool_size':     args.pool_size,
            'initial_frames': args.initial_frames,
            'max_iterations': args.max_iterations,
        },
        checkpoint_every=args.checkpoint_every,
    )


def _model_id_for(backbone: str) -> str:
    from util.backbones import DEFAULT_MODELS, _configs
    return (_configs.get(backbone, {}) or {}).get('model_id') or DEFAULT_MODELS.get(backbone, backbone)


# ---------------------------------------------------------------------------
#                     The agentic loop (per question)
# ---------------------------------------------------------------------------
def _agentic_loop(
    *,
    question_row: dict,
    movie_title: str,
    all_frames: list[Image.Image],
    all_timestamps: np.ndarray,
    clip_embs: np.ndarray,
    args: argparse.Namespace,
    clip_model: Any,
    clip_processor: Any,
) -> tuple[dict, int, int]:

    backbone = args.backbone
    selected_indices = np.linspace(
        0, len(all_frames) - 1, args.initial_frames, dtype=int
    ).tolist()
    final_preds: dict = {}
    raw_answer_text = ''

    for iteration in range(args.max_iterations):
        n_iterations = iteration + 1
        selected_images = [all_frames[i] for i in selected_indices]
        frame_labels = [f'Frame {i} (t={all_timestamps[i]:.1f}s)' for i in selected_indices]

        try:
            captions = _caption_frames(selected_images, frame_labels, backbone, args)
        except Exception as e:
            print(f'  [caption err] qid={question_row["qid"]} iter={iteration}: {type(e).__name__}: {e!r}')
            time.sleep(2)
            break

        try:
            final_preds, raw_answer_text = _answer_mcq(
                captions, question_row, movie_title, backbone, args,
            )
        except Exception as e:
            print(f'  [answer err] qid={question_row["qid"]} iter={iteration}: {type(e).__name__}: {e!r}')
            time.sleep(2)
            break

        try:
            confidence = _self_evaluate(captions, question_row, raw_answer_text, backbone, args)
        except Exception as e:
            print(f'  [confidence err] qid={question_row["qid"]} iter={iteration}: {type(e).__name__}: {e!r}')
            confidence = args.confidence_threshold  # assume sufficient so we stop cleanly
            time.sleep(2)

        if confidence >= args.confidence_threshold:
            break

        try:
            segment_descs, segment_map = _identify_segments(
                captions, question_row, selected_indices, len(all_frames), backbone, args,
            )
        except Exception as e:
            print(f'  [segment err] qid={question_row["qid"]} iter={iteration}: {type(e).__name__}: {e!r}')
            break
        if not segment_descs:
            break

        try:
            new_indices = _retrieve_frames(
                segment_descs, segment_map, clip_embs, selected_indices,
                args.frames_per_retrieval, clip_model, clip_processor,
            )
        except Exception as e:
            print(f'  [clip retrieval err] qid={question_row["qid"]} iter={iteration}: '
                  f'{type(e).__name__}: {e!r}')
            break
        if not new_indices:
            break

        selected_indices = sorted(set(selected_indices + new_indices))

    return final_preds, n_iterations, len(selected_indices)


def _caption_frames(images, frame_labels, backbone: str, args) -> dict:
    prompt = (
        'Describe each of the following video frames in 1-2 sentences. '
        'Focus on visible objects, people, actions, and scene context. '
        'Return a JSON object mapping frame labels to their descriptions.\n\n'
        f'Frame labels: {", ".join(frame_labels)}'
    )
    raw = call_backbone(
        backbone, prompt, images=images,
        system_instruction='You caption video frames. Output valid JSON.',
        json_mode=True,
        max_output_tokens=args.max_output_tokens,
        temperature=args.temperature,
    )
    parsed = parse_json_flexible(raw)
    if isinstance(parsed, dict):
        return parsed
    # Fallback: attribute raw text to every frame
    return {label: (raw or '') for label in frame_labels}


def _build_q_json(row: dict) -> str:
    return json.dumps({
        'id': row['qid'],
        'question': row['question'] + ' Select the single best choice.',
        'A': row['option_A'], 'B': row['option_B'],
        'C': row['option_C'], 'D': row['option_D'],
    })


def _answer_mcq(captions: dict, row: dict, movie_title: str, backbone: str, args) -> tuple[dict, str]:
    caption_text = json.dumps(captions, indent=2)
    question_json = _build_q_json(row)
    movie_ctx = f' from the movie "{movie_title}"' if movie_title else ''
    prompt = (
        f'Given the following descriptions of sampled frames from a video{movie_ctx}:\n\n'
        f'{caption_text}\n\n'
        f'Answer the following question:\n\n{question_json}\n\n'
        'Rules:\n'
        '- You MUST answer with one of the choices. If uncertain, commit to your '
        'best guess based on the descriptions and your prior knowledge.\n'
        '- Do NOT refuse, do NOT ask for more frames, do NOT include commentary -— '
        'output ONLY the JSON below.\n\n'
        'Output ONLY valid JSON in this format:\n'
        '[{"id": id, "answers": ["X"]}]        // single-choice\n'
        '[{"id": id, "answers": ["X", "Y"]}]   // multi-answer'
    )
    raw = call_backbone(
        backbone, prompt,
        system_instruction='You answer movie MCQs from frame descriptions. Output valid JSON.',
        json_mode=True,
        max_output_tokens=args.max_output_tokens,
        temperature=args.temperature,
    )
    parsed_text = parse_mcq_json(raw, qid=row['qid'])
    preds = {norm_qid(row['qid']): parsed_text} if parsed_text else {}
    return preds, raw


def _self_evaluate(captions: dict, row: dict, answer_text: str, backbone: str, args) -> int:
    prompt = (
        'Please assess the confidence level in the decision-making process.\n\n'
        f'Frame descriptions:\n{json.dumps(captions, indent=2)}\n\n'
        f'Question:\n{_build_q_json(row)}\n\n'
        f'The decision/answer is:\n{answer_text}\n\n'
        'Criteria:\n'
        '1 = Insufficient: too little info for a reasonable conclusion.\n'
        '2 = Partial: supports an informed guess only.\n'
        '3 = Sufficient: supports a well-informed decision.\n\n'
        'Return JSON: {"confidence": 1|2|3}'
    )
    raw = call_backbone(
        backbone, prompt,
        system_instruction='You self-assess answer confidence. Output valid JSON.',
        json_mode=True,
        max_output_tokens=args.max_output_tokens,
        temperature=args.temperature,
    )
    parsed = parse_json_flexible(raw)
    try:
        conf = int(parsed['confidence'])
        if conf in (1, 2, 3):
            return conf
    except (TypeError, ValueError, KeyError):
        pass
    return 1   # lowest, forces another iteration on parse failure


def _identify_segments(
    captions: dict, row: dict, selected_indices: list[int], n_total_frames: int,
    backbone: str, args,
) -> tuple[list[dict], dict[int, str]]:
    sorted_idx = sorted(selected_indices)
    segment_des = {
        i + 1: f'{sorted_idx[i]}-{sorted_idx[i + 1]}'
        for i in range(len(sorted_idx) - 1)
    }
    example = {
        'frame_descriptions': [
            {'segment_id': '1', 'duration': 'xxx - xxx', 'description': 'frame of xxx'},
            {'segment_id': '2', 'duration': 'xxx - xxx', 'description': 'frame of xxx'},
        ]
    }
    prompt = (
        f'Given a video with {n_total_frames} frames. '
        f'Descriptions of sampled frames:\n\n{json.dumps(captions, indent=2)}\n\n'
        f'To answer:\n{_build_q_json(row)}\n\n'
        f'Initial frames are insufficient. Candidate segments (from gaps): '
        f'{json.dumps(segment_des)}\n\n'
        f'For each segment likely to contain relevant info, give a concise '
        f'description of what to look for.\n\n'
        f'Return JSON with segment_id (must be ≤ {len(segment_des)}), '
        f'duration, description:\n{json.dumps(example, indent=2)}'
    )
    raw = call_backbone(
        backbone, prompt,
        system_instruction='You direct visual search for MCQs. Output valid JSON.',
        json_mode=True,
        max_output_tokens=args.max_output_tokens,
        temperature=args.temperature,
    )
    parsed = parse_json_flexible(raw)
    if isinstance(parsed, dict) and 'frame_descriptions' in parsed:
        return parsed['frame_descriptions'], segment_des
    if isinstance(parsed, list):
        return parsed, segment_des
    return [], segment_des


def _retrieve_frames(
    descriptions: list[dict], segment_des: dict[int, str],
    clip_embeddings: np.ndarray, existing_indices: list[int],
    frames_per_retrieval: int, clip_model, clip_processor,
) -> list[int]:
    new_indices: list[int] = []
    for item in descriptions:
        try:
            seg_id = int(item.get('segment_id', 1))
        except (TypeError, ValueError):
            continue
        if seg_id not in segment_des:
            continue
        seg_range = segment_des[seg_id]
        try:
            seg_start, seg_end = [int(x.strip()) for x in seg_range.split('-')]
        except ValueError:
            continue
        if seg_start >= len(clip_embeddings) or seg_end > len(clip_embeddings):
            continue
        seg_embs = clip_embeddings[seg_start:seg_end + 1]
        if len(seg_embs) == 0:
            continue

        query = item.get('description', '')
        if not query:
            continue
        top_k = min(2, len(seg_embs))
        idx_in_seg, _ = retrieve_frames_by_text(
            query, seg_embs,
            clip_model=clip_model, clip_processor=clip_processor,
            top_k=top_k,
        )
        for local_idx in idx_in_seg:
            global_idx = seg_start + local_idx
            if global_idx not in existing_indices and global_idx not in new_indices:
                new_indices.append(global_idx)
    return new_indices[:frames_per_retrieval]


if __name__ == '__main__':
    main()