#!/usr/bin/env python
"""
End-to-end MCQ evaluation with Hour-LLaVA (long-video VLM).

Usage:
    python run.py --frame-budgets 128 256
    python run.py --model-id jylins/hourllava-qwen25-3b-s3
"""
from __future__ import annotations

import argparse
import re
import sys
import types
from pathlib import Path

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

_DEPS = _PROJECT / 'dependencies' / 'hourllava'
sys.path.insert(0, str(_DEPS))

import numpy as np
import torch

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

from decord import VideoReader, cpu


EXPERIMENT_NAME = 'hourllava'


def _sample_frames(vpath: Path, n: int) -> np.ndarray:
    vr = VideoReader(str(vpath), ctx=cpu(0))
    idx = np.linspace(0, len(vr) - 1, n, dtype=int)
    return vr.get_batch(idx).asnumpy()


def _patch_dist_barrier_noop() -> None:
    """dist.barrier() assumes DDP; no-op when distributed isn't initialized."""
    import torch.distributed as _dist
    if not getattr(_dist, '_chnb_patched', False):
        _orig = _dist.barrier
        def _noop(*a, **kw):
            if not _dist.is_initialized():
                return
            return _orig(*a, **kw)
        _dist.barrier = _noop
        _dist._chnb_patched = True


def load_hourllava(model_id: str, device_map: str):
    """Bypass hourllava.model.builder (buggy dispatch) and load directly."""
    _patch_dist_barrier_noop()

    from transformers import AutoTokenizer
    from hourllava.model.language_model.hour_llava_qwen import (
        HourLlavaQwenForCausalLM,
    )
    from hourllava.util.constants import DEFAULT_IMAGE_PATCH_TOKEN, IGNORE_INDEX

    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
    model = HourLlavaQwenForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map=device_map,
        attn_implementation='flash_attention_2',
        low_cpu_mem_usage=True,
    )

    vision_tower = model.get_vision_tower()
    if not vision_tower.is_loaded:
        vision_tower.load_model(device_map=device_map)
    vision_tower.to(dtype=model.dtype)
    image_processor = vision_tower.image_processor

    # Image patch token
    if getattr(model.config, 'mm_use_im_patch_token', True):
        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
        model.resize_token_embeddings(len(tokenizer))

    # Missing runtime attrs on inner model
    inner = model.get_model()
    for attr, default in [('spatial_sampling_type', 'uniform'),
                          ('temporal_sampling_type', 'uniform')]:
        if not hasattr(inner, attr):
            setattr(inner, attr, getattr(model.config, attr, default))

    # Instance-level patch for prepare_inputs_labels_for_multimodal (labels=None bug)
    _orig_prepare = model.prepare_inputs_labels_for_multimodal
    def _patched(self, input_ids, position_ids, attention_mask, past_key_values,
                 labels, images, modalities=['image'], image_sizes=None):
        if labels is None:
            labels = torch.full_like(input_ids, IGNORE_INDEX)
        return _orig_prepare(input_ids, position_ids, attention_mask,
                             past_key_values, labels, images, modalities,
                             image_sizes=image_sizes)
    model.prepare_inputs_labels_for_multimodal = types.MethodType(_patched, model)

    return tokenizer, model, image_processor


def run_inference(frames, prompt_text: str, *, tokenizer, model, image_processor,
                  max_new_tokens: int) -> str:
    from hourllava.util.constants import (
        DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX,
    )
    from hourllava.util.conversation import conv_templates
    from hourllava.util.mm_utils import tokenizer_image_token

    # Prepend the <image> placeholder so `tokenizer_image_token` actually
    # inserts an IMAGE_TOKEN_INDEX into input_ids. Without this, num_images_list
    # is [0] but modalities=['video'] (length 1), and the assert at
    # hour_llava_arch.py:326 fires (`1 != 0`).
    prompt_with_image = f'{DEFAULT_IMAGE_TOKEN}\n{prompt_text}'

    conv = conv_templates['qwen_1_5'].copy()
    conv.append_message(conv.roles[0], prompt_with_image)
    conv.append_message(conv.roles[1], None)
    full_prompt = conv.get_prompt()

    input_ids = tokenizer_image_token(
        full_prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt',
    ).unsqueeze(0).to(model.device)

    pixel_values = image_processor.preprocess(frames, return_tensors='pt')['pixel_values'].to(
        model.device, dtype=model.dtype,
    )

    # Hour-LLaVA wants a list of 4D tensors (NOT a 4D tensor directly)
    with torch.no_grad():
        output_ids = model.generate(
            inputs=input_ids,
            images=[pixel_values],
            image_sizes=[(frames[0].shape[1], frames[0].shape[0])],
            modalities=['video'],
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=1.0,
        )

    # Hour-LLaVA's generate() consumes inputs_embeds internally, so output_ids
    # contains only the generated tokens (no input prefix).
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--frame-budgets', type=int, nargs='+', default=[0, 64, 128, 256])
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--model-id', default=DEFAULT_MODELS['hourllava'],
                    help='Try jylins/hourllava-qwen25-3b-s3 for a smaller variant')
    ap.add_argument('--max-new-tokens', type=int, default=1024)
    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 {args.model_id}...')
    tokenizer, model, image_processor = load_hourllava(args.model_id, args.device_map)
    print(f'Ready: dtype={model.dtype}, vision_tower.dtype={next(model.get_vision_tower().parameters()).dtype}')

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


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

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

    # Track first-failure traceback so a recurring AssertionError shows its
    # actual location (which assert, which file) instead of just the empty
    # `AssertionError()` repr that's all you get from `e!r`.
    _tb_shown = {'text': False, 'video': False}
    import traceback as _tb_mod

    def call_fn(row, ctx):
        frames = ctx.get('frames')
        if num_frames > 0 and frames is None:
            return ''

        if num_frames == 0:
            prompt = mcq_prompt_json(row, header=HEADER_TEXT_ONLY)
            try:
                messages = [
                    {'role': 'system', 'content': 'You answer movie MCQs. Output only valid JSON.'},
                    {'role': 'user', 'content': prompt},
                ]
                text = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True,
                )
                input_ids = tokenizer(text, return_tensors='pt').input_ids.to(model.device)
                with torch.no_grad():
                    output_ids = model.generate(
                        inputs=input_ids,
                        max_new_tokens=args.max_new_tokens,
                        do_sample=False,
                        temperature=1.0,
                    )
                # Strip the input prefix (Hour-LLaVA's inputs_embeds path may or
                # may not include it; slice defensively).
                gen = output_ids[0]
                if gen.shape[0] > input_ids.shape[1] and torch.equal(
                    gen[: input_ids.shape[1]], input_ids[0]
                ):
                    gen = gen[input_ids.shape[1]:]
                raw = tokenizer.decode(gen, skip_special_tokens=True)
            except Exception as e:
                print(f'  text-only error {row["qid"]}: {type(e).__name__}: {e!r}')
                if not _tb_shown['text']:
                    _tb_mod.print_exc()
                    _tb_shown['text'] = True
                return ''
            return parse_mcq_json(raw, qid=row['qid'])

        prompt = mcq_prompt_json(row)
        try:
            raw = run_inference(
                frames, prompt,
                tokenizer=tokenizer, model=model, image_processor=image_processor,
                max_new_tokens=args.max_new_tokens,
            )
        except Exception as e:
            print(f'  inference error {row["qid"]}: {type(e).__name__}: {e!r}')
            if not _tb_shown['video']:
                _tb_mod.print_exc()
                _tb_shown['video'] = True
            return ''
        return parse_mcq_json(raw, qid=row['qid'])

    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, args.model_id),
        experiment_name=f'{EXPERIMENT_NAME} f={num_frames}',
        extra_row_fields={
            'num_frames': num_frames,
            'model':      args.model_id,
            'backbone':   'hourllava',
            '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()