"""
Shared per-question evaluation loop.

A per-model script reduces to: load benchmark, define `call_fn`, call
`run_sweep(...)`. The loop is resumable, checkpoints periodically, drops
empty-pred and multi-letter rows on load so they get retried, and prints
overall accuracy with a 95% bootstrap CI stratified by movie at the end.

Example:

    from util.benchmark import load_benchmark
    from util.paths import BENCHMARK_PATH
    from util.eval_loop import run_sweep

    df = load_benchmark(BENCHMARK_PATH)

    def call_fn(row, ctx):
        return my_model_answer(row['question'], ctx['frames'])

    def build_movie_ctx(vid, movie_df):
        return {'frames': extract_frames(vid)}

    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        results_path='output/results/my_exp/results_f128.csv',
        extra_row_fields={'num_frames': 128, 'model': 'my-model'},
    )
"""
from __future__ import annotations

import shutil
import time
import traceback as _tb
from datetime import datetime
from pathlib import Path
from typing import Callable

import pandas as pd
from tqdm.auto import tqdm


_CANONICAL_COLUMNS = [
    'qid', 'video_file_id', 'category', 'q_type',
    'gold', 'pred', 'latency_s',
]


def _load_existing(results_path: Path) -> list[dict]:
    """Load CSV if it exists, drop empty / multi-letter preds, back it up.

    Returns the records list (one dict per row).
    """
    if not results_path.exists():
        return []

    # Backup
    ts = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_path = results_path.with_name(results_path.stem + f'_backup_{ts}.csv')
    shutil.copy2(results_path, backup_path)
    print(f'  Backed up -> {backup_path}')

    if results_path.stat().st_size == 0:
        print('  Existing CSV empty —- starting fresh')
        return []

    try:
        df = pd.read_csv(results_path)
    except pd.errors.EmptyDataError:
        print('  Existing CSV unparseable —- starting fresh')
        return []

    # Drop empty preds (budget-bound or transient API failures — retry on
    # resume) and multi-letter preds ('A,C') that don't score under SC.
    if 'pred' in df.columns and len(df):
        pred_clean = df['pred'].fillna('').astype(str).str.strip().str.upper()
        is_sc_valid = pred_clean.isin({'A', 'B', 'C', 'D'})
        is_empty    = pred_clean == ''
        n_empty = int(is_empty.sum())
        n_multi = int((~is_sc_valid & ~is_empty).sum())
        df = df[is_sc_valid].copy()
        if n_empty:
            print(f'  Dropped {n_empty} empty-pred rows — will retry on this resume')
        if n_multi:
            print(f'  Dropped {n_multi} non-SC-valid pred rows (multi-letter / malformed)')

    # Dedup on qid (last write wins).
    if 'qid' in df.columns and len(df):
        before = len(df)
        df = df.drop_duplicates(subset='qid', keep='last').reset_index(drop=True)
        dropped = before - len(df)
        if dropped:
            print(f'  Deduped {dropped} duplicate-qid rows (kept last write)')

    return df.to_dict('records')


def _emit_row(row: dict, pred: str, extra: dict | None,
              latency_s: float | None = None) -> dict:
    """Build one output row with canonical columns + any extras (condition tags)."""
    out = {k: row.get(k, '') for k in _CANONICAL_COLUMNS}
    out['pred'] = pred
    out['latency_s'] = float(latency_s) if latency_s is not None else 0.0
    if extra:
        for k, v in extra.items():
            if k != 'latency_s':
                out[k] = v
    return out


def _print_final_summary(out_df: pd.DataFrame) -> None:
    """Print overall accuracy with a 95% bootstrap CI stratified by movie."""
    if len(out_df) == 0:
        return
    try:
        from util.metrics import bootstrap_ci, fmt
    except Exception as e:
        print(f'  (Skipping accuracy summary: {type(e).__name__}: {e})')
        return

    gold = out_df['gold'].fillna('').astype(str).str.strip().str.upper()
    pred = out_df['pred'].fillna('').astype(str).str.strip().str.upper()
    correct = (gold == pred) & gold.isin(['A', 'B', 'C', 'D'])
    strata = out_df['video_file_id'].fillna('').astype(str)

    stat = bootstrap_ci(correct.values.astype(float), strata.values)
    n = len(out_df)
    n_correct = int(correct.sum())
    print()
    print(f'  Accuracy: {fmt(stat)}  ({n_correct}/{n} correct, '
          f'95% CI stratified by movie, 10K bootstrap iterations)')


def run_sweep(
    df: pd.DataFrame,
    *,
    call_fn: Callable[[dict, dict], str],
    results_path: str | Path,
    experiment_name: str = 'eval',
    extra_row_fields: dict | None = None,
    build_movie_ctx: Callable[[str, pd.DataFrame], dict] | None = None,
    teardown_movie_ctx: Callable[[dict], None] | None = None,
    checkpoint_every: int = 50,
    skip_if_no_ctx: bool = True,
) -> pd.DataFrame:
    """
    Per-question sweep with per-movie context reuse.

    Args:
        df: benchmark DataFrame (must include 'qid' and 'video_file_id').
        call_fn: (row_dict, movie_ctx_dict) -> pred_str. Return '' to trigger retry next run.
        results_path: where to write CSV (parent dirs created automatically).
        experiment_name: used in log prefixes.
        extra_row_fields: dict merged into every output row (e.g. {'num_frames': 128}).
        build_movie_ctx: optional (vid, movie_df) -> dict — built once per movie and passed
            into every call_fn for that movie. Use for frame caching, video extraction, etc.
        teardown_movie_ctx: optional (movie_ctx_dict) -> None — called after last question
            in a movie. Use to free GPU memory, close file handles, etc.
        checkpoint_every: save CSV every N new rows.
        skip_if_no_ctx: if build_movie_ctx returns None/empty dict, skip the movie's
            remaining questions (they'll retry next run). Default True.

    Returns:
        Final pandas DataFrame of all results (written to `results_path`).
    """
    results_path = Path(results_path)
    results_path.parent.mkdir(parents=True, exist_ok=True)

    print(f'\n=== {experiment_name} ===')
    print(f'Results -> {results_path}')

    # Load + backup + drop empties
    all_results = _load_existing(results_path)
    existing_qids = set(str(r.get('qid', '')) for r in all_results)

    # Figure out remaining questions
    remaining = [row.to_dict() for _, row in df.iterrows()
                 if str(row['qid']) not in existing_qids]
    print(f'{len(all_results)} existing rows, {len(remaining)} questions remaining')
    if not remaining:
        out_df = pd.DataFrame(all_results)
        _print_final_summary(out_df)
        return out_df

    _hist = [float(r.get('latency_s', 0.0) or 0.0) for r in all_results]
    _hist = [x for x in _hist if x > 0]
    if _hist:
        _hist_sorted = sorted(_hist)
        _med = _hist_sorted[len(_hist_sorted) // 2]
        _mean = sum(_hist) / len(_hist)
        eta_sec = _med * len(remaining)
        if eta_sec < 60:
            eta_str = f'{eta_sec:.0f}s'
        elif eta_sec < 3600:
            eta_str = f'{eta_sec/60:.1f}min'
        else:
            eta_str = f'{eta_sec/3600:.1f}h'
        print(f'  Historical latency (n={len(_hist)} prior calls): '
              f'median={_med:.2f}s, mean={_mean:.2f}s')
        print(f'  ETA for remaining {len(remaining)} questions: ~{eta_str}')
    else:
        print(f'  No historical latency data; tqdm ETA will populate after the '
              f'first few calls.')

    # Group remaining by movie while preserving original order
    movies_in_order: list[str] = []
    by_movie: dict[str, list[dict]] = {}
    for r in remaining:
        vid = r['video_file_id']
        if vid not in by_movie:
            movies_in_order.append(vid)
            by_movie[vid] = []
        by_movie[vid].append(r)

    pbar = tqdm(total=len(remaining), desc=experiment_name)
    _traceback_shown = False

    sweep_t0 = time.perf_counter()
    n_before = len(all_results)

    try:
        _movie_loop_body(
            movies_in_order, by_movie, df, build_movie_ctx, teardown_movie_ctx,
            call_fn, extra_row_fields, all_results, existing_qids,
            results_path, checkpoint_every, skip_if_no_ctx, pbar, _traceback_shown,
        )
    except BaseException as e:
        cls = type(e).__name__
        if cls == 'FatalBackboneError':
            print(f'\nFATAL: {e}')
            print('Aborting sweep. Partial results saved.')
        elif cls == 'KeyboardInterrupt':
            print('\nInterrupted. Partial results saved.')
        pbar.close()
        if all_results:
            pd.DataFrame(all_results).to_csv(results_path, index=False)
        sweep_dt = time.perf_counter() - sweep_t0
        n_new = len(all_results) - n_before
        print(f'  Partial wall time: {sweep_dt:.1f}s for {n_new} new rows.')
        raise

    pbar.close()
    out_df = pd.DataFrame(all_results)
    out_df.to_csv(results_path, index=False)
    sweep_dt = time.perf_counter() - sweep_t0
    n_new = len(all_results) - n_before
    new_rows = all_results[n_before:]
    sum_call_s = sum(float(r.get('latency_s', 0.0) or 0.0) for r in new_rows)
    print(f'  Done: {len(out_df)} rows -> {results_path}')
    if n_new > 0:
        mean_call = sum_call_s / n_new
        overhead = max(0.0, sweep_dt - sum_call_s)
        print(f'  Wall time: {sweep_dt:.1f}s ({n_new} new rows; '
              f'mean call {mean_call:.2f}s, '
              f'sum call {sum_call_s:.1f}s, '
              f'non-call overhead {overhead:.1f}s — frame extraction, '
              f'caption I/O, ctx teardown, etc.)')
    else:
        print(f'  Wall time: {sweep_dt:.1f}s (no new rows produced)')

    _print_final_summary(out_df)

    return out_df


def _movie_loop_body(movies_in_order, by_movie, df, build_movie_ctx,
                     teardown_movie_ctx, call_fn, extra_row_fields,
                     all_results, existing_qids, results_path,
                     checkpoint_every, skip_if_no_ctx, pbar, _traceback_shown):
    """Inner loop extracted so the outer run_sweep can wrap it for clean abort."""
    n_call_err = 0
    n_empty    = 0
    for vid in movies_in_order:
        movie_rows = by_movie[vid]
        movie_df = df[df['video_file_id'] == vid]

        # Build per-movie ctx
        ctx: dict = {}
        if build_movie_ctx is not None:
            try:
                ctx = build_movie_ctx(vid, movie_df) or {}
            except Exception as e:
                print(f'  [ctx build] {vid}: {type(e).__name__}: {e!r}')
                if not _traceback_shown:
                    _tb.print_exc()
                    _traceback_shown = True
                ctx = {}

        if skip_if_no_ctx and build_movie_ctx is not None and not ctx:
            for _ in movie_rows:
                pbar.update(1)
            continue

        for row in movie_rows:
            call_t0 = time.perf_counter()
            try:
                pred = call_fn(row, ctx)
            except Exception as e:
                # FatalBackboneError is BaseException, propagates past this catch.
                print(f'  [call_fn] {vid} qid={row.get("qid")}: {type(e).__name__}: {e!r}')
                if not _traceback_shown:
                    _tb.print_exc()
                    _traceback_shown = True
                n_call_err += 1
                pred = ''
            call_dt = time.perf_counter() - call_t0

            pred_str = str(pred).strip() if pred is not None else ''
            all_results.append(
                _emit_row(row, pred_str, extra_row_fields, latency_s=call_dt)
            )
            existing_qids.add(str(row['qid']))
            if len(all_results) % checkpoint_every == 0:
                pd.DataFrame(all_results).to_csv(results_path, index=False)
            if not pred_str:
                n_empty += 1
                if n_empty <= 3:
                    print(f'  [empty pred] {vid} qid={row.get("qid")} '
                          f'(parser returned ""; recorded as terminal failure)')
            pbar.update(1)
            pbar.set_postfix({'ok': len(all_results),
                              'empty': n_empty, 'err': n_call_err})

        if teardown_movie_ctx is not None:
            try:
                teardown_movie_ctx(ctx)
            except Exception as e:
                print(f'  [ctx teardown] {vid}: {type(e).__name__}: {e!r}')

    total_processed = len(all_results) + n_empty + n_call_err
    if (n_empty or n_call_err) and total_processed > 0:
        rate = (n_empty + n_call_err) / total_processed
        print(f'  Sweep summary: {len(all_results)} written, '
              f'{n_empty} empty, {n_call_err} errors '
              f'({rate:.0%} failure rate)')
        if rate >= 0.5:
            print(f'  Failure rate >= 50%. Likely causes: model context-window '
                  f'exceeded, parser too strict, or upstream API rejecting requests. '
                  f'Inspect a raw response before re-running.')