#!/usr/bin/env python3
"""
Run one CHNB question through the public end-to-end model runners.

This script intentionally shells out to each experiment's real `run.py` with a
one-row benchmark file and isolated CHNB_OUTPUT_ROOT. That keeps one-question
inference aligned with full-sweep behavior instead of maintaining a separate
inference implementation.
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import shlex
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable


PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.dont_write_bytecode = True
sys.path.insert(0, str(PROJECT_ROOT))

from util.paths import BENCHMARK_PATH, OUTPUT_ROOT


DEFAULT_BENCHMARK = BENCHMARK_PATH


@dataclass(frozen=True)
class Runner:
    experiment: str
    workdir: Path
    frame_budget_arg: bool = True
    extra_args: tuple[str, ...] = ()


RUNNERS: dict[str, Runner] = {
    'gemini3flash': Runner(
        experiment='gemini3flash',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'gemini3flash',
    ),
    'gpt5mini': Runner(
        experiment='gpt5mini',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'gpt5mini',
    ),
    'claude45haiku': Runner(
        experiment='claude45haiku',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'claude45haiku',
    ),
    'qwen3vl': Runner(
        experiment='qwen3vl',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'qwen3vl',
    ),
    'glm4v': Runner(
        experiment='glm4v',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'glm4v',
    ),
    'llava_next_video': Runner(
        experiment='llava_next_video',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'llava_next_video',
    ),
    'hourllava': Runner(
        experiment='hourllava',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'hourllava',
    ),
    'vamba': Runner(
        experiment='vamba',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'vamba',
    ),
    'videochat': Runner(
        experiment='videochat',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'videochat',
    ),
    # Full-video Gemini is not in the default "all" set because it requires
    # movies in GCS rather than local files / pre-extracted frames.
    'gemini3flash_video': Runner(
        experiment='gemini3flash_video',
        workdir=PROJECT_ROOT / 'experiments' / 'end-to-end' / 'gemini3flash_video',
        frame_budget_arg=False,
    ),
}

DEFAULT_MODELS = (
    'gemini3flash',
    'gpt5mini',
    'claude45haiku',
    'qwen3vl',
    'glm4v',
    'llava_next_video',
    'hourllava',
    'vamba',
    'videochat',
)


def _load_jsonl(path: Path) -> list[dict]:
    rows: list[dict] = []
    with path.open(encoding='utf-8') as f:
        for line in f:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def _select_question(rows: list[dict], *, qid: str | None, index: int) -> dict:
    if not rows:
        raise SystemExit('Benchmark file has no rows.')
    if qid is not None:
        for row in rows:
            if str(row.get('qid')) == str(qid):
                return row
        raise SystemExit(f'No question with qid={qid!r}.')
    if index < 0 or index >= len(rows):
        raise SystemExit(f'Question index out of range: {index}; benchmark has {len(rows)} rows.')
    return rows[index]


def _expand_models(tokens: Iterable[str]) -> list[str]:
    out: list[str] = []
    for token in tokens:
        if token == 'all':
            out.extend(DEFAULT_MODELS)
        elif token == 'all-plus-video':
            out.extend(DEFAULT_MODELS)
            out.append('gemini3flash_video')
        else:
            out.append(token)

    unknown = [m for m in out if m not in RUNNERS]
    if unknown:
        raise SystemExit(
            f'Unknown model runner(s): {", ".join(unknown)}\n'
            f'Known runners: {", ".join(RUNNERS)}'
        )

    deduped: list[str] = []
    for model in out:
        if model not in deduped:
            deduped.append(model)
    return deduped


def _parse_runner_args(values: list[str]) -> dict[str, list[str]]:
    parsed: dict[str, list[str]] = {}
    for value in values:
        if '=' not in value:
            raise SystemExit('--runner-arg must look like model="--flag value"')
        model, raw_args = value.split('=', 1)
        model = model.strip()
        if model not in RUNNERS:
            raise SystemExit(f'Unknown --runner-arg target: {model}')
        parsed.setdefault(model, []).extend(shlex.split(raw_args))
    return parsed


def _write_one_question(row: dict, out_dir: Path) -> Path:
    path = out_dir / f'question_{row["qid"]}.jsonl'
    path.write_text(json.dumps(row, ensure_ascii=False) + '\n', encoding='utf-8')
    return path


def _run_and_log(name: str, cmd: list[str], cwd: Path, env: dict[str, str],
                 log_path: Path, *, quiet: bool) -> int:
    with log_path.open('w', encoding='utf-8') as log:
        log.write('$ ' + ' '.join(shlex.quote(x) for x in cmd) + '\n\n')
        proc = subprocess.Popen(
            cmd,
            cwd=str(cwd),
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
        assert proc.stdout is not None
        for line in proc.stdout:
            log.write(line)
            if not quiet:
                print(f'[{name}] {line}', end='')
        return proc.wait()


def _latest_result_row(output_root: Path, experiment: str, qid: str) -> dict | None:
    result_dir = output_root / 'results' / experiment
    if not result_dir.exists():
        return None
    csv_paths = sorted(result_dir.glob('*.csv'), key=lambda p: p.stat().st_mtime, reverse=True)
    for path in csv_paths:
        try:
            with path.open(encoding='utf-8', newline='') as f:
                for row in csv.DictReader(f):
                    if str(row.get('qid')) == str(qid):
                        row['_results_csv'] = str(path)
                        return row
        except Exception:
            continue
    return None


def _write_summary(rows: list[dict], out_dir: Path) -> None:
    summary_json = out_dir / 'summary.json'
    summary_csv = out_dir / 'summary.csv'
    accuracy_json = out_dir / 'accuracy.json'
    summary_json.write_text(json.dumps(rows, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
    accuracy_json.write_text(
        json.dumps(_accuracy_summary(rows), indent=2, ensure_ascii=False) + '\n',
        encoding='utf-8',
    )

    columns = [
        'runner', 'status', 'qid', 'video_file_id', 'model', 'backbone',
        'num_frames', 'gold', 'pred', 'correct', 'latency_s',
        'results_csv', 'log',
    ]
    with summary_csv.open('w', encoding='utf-8', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=columns)
        writer.writeheader()
        for row in rows:
            writer.writerow({k: row.get(k, '') for k in columns})


def _is_correct(row: dict) -> bool:
    value = row.get('correct')
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() == 'true'


def _accuracy_summary(rows: list[dict]) -> dict:
    answered = [row for row in rows if row.get('pred')]
    correct = sum(1 for row in answered if _is_correct(row))
    errors = [row for row in rows if str(row.get('status', '')).startswith('error:')]
    no_result = [row for row in rows if row.get('status') == 'no_result']
    dry_run = [row for row in rows if row.get('status') == 'dry_run']
    accuracy = (correct / len(answered)) if answered else None
    return {
        'correct': correct,
        'answered': len(answered),
        'total_runners': len(rows),
        'accuracy': accuracy,
        'errors': len(errors),
        'no_result': len(no_result),
        'dry_run': len(dry_run),
    }


def _print_accuracy_summary(rows: list[dict]) -> None:
    stats = _accuracy_summary(rows)
    if stats['answered']:
        pct = 100.0 * stats['accuracy']
        print(
            f'One-question accuracy: {stats["correct"]}/{stats["answered"]} '
            f'= {pct:.1f}% over runners that produced a prediction.'
        )
    elif stats['dry_run']:
        print('One-question accuracy: n/a (dry-run; no model predictions produced).')
    else:
        print('One-question accuracy: n/a (no runner produced a prediction).')

    if stats['errors'] or stats['no_result']:
        print(
            f'Run issues: errors={stats["errors"]}, '
            f'no_result={stats["no_result"]}, total_runners={stats["total_runners"]}'
        )


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--benchmark', type=Path, default=DEFAULT_BENCHMARK)
    ap.add_argument('--qid', default=None, help='Question id to run. Defaults to --index.')
    ap.add_argument('--index', type=int, default=0, help='0-based question index when --qid is omitted.')
    ap.add_argument('--frames', type=int, default=64,
                    help='Frame budget passed to frame/video runners. Default 64 fits every public runner.')
    ap.add_argument('--models', nargs='+', default=['all'],
                    help='Model runners to execute. Use "all" or "all-plus-video".')
    ap.add_argument('--output-dir', type=Path, default=None,
                    help='Directory for logs and summaries. Default: output/proc_one/<timestamp>_qid_<qid>.')
    ap.add_argument('--runner-arg', action='append', default=[],
                    help='Extra args for one runner, e.g. gpt5mini="--model gpt-5". Repeatable.')
    ap.add_argument('--fail-fast', action='store_true')
    ap.add_argument('--require-all-ok', action='store_true',
                    help='After trying all selected runners, exit nonzero if any runner failed '
                         'or did not produce a result.')
    ap.add_argument('--quiet', action='store_true', help='Write runner output to logs without streaming it.')
    ap.add_argument('--dry-run', action='store_true', help='Print commands without executing them.')
    args = ap.parse_args()

    rows = _load_jsonl(args.benchmark)
    question = _select_question(rows, qid=args.qid, index=args.index)
    qid = str(question['qid'])
    models = _expand_models(args.models)
    runner_args = _parse_runner_args(args.runner_arg)

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    out_dir = args.output_dir or OUTPUT_ROOT / 'proc_one' / f'{timestamp}_qid_{qid}'
    out_dir = out_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    one_question_path = _write_one_question(question, out_dir)
    output_root = out_dir / 'runner_outputs'
    output_root.mkdir(parents=True, exist_ok=True)

    print(f'Question qid={qid}  video={question["video_file_id"]}')
    print(f'Gold: {question.get("answer", "")}')
    print(f'Frame budget: {args.frames}')
    print(f'Output: {out_dir}')
    print(f'Runners: {", ".join(models)}')
    print()

    env = os.environ.copy()
    env.setdefault('CHNB_PROJECT_ROOT', str(PROJECT_ROOT))
    env.update({
        'CHNB_OUTPUT_ROOT': str(output_root),
        'PYTHONDONTWRITEBYTECODE': '1',
    })

    summary: list[dict] = []
    for name in models:
        runner = RUNNERS[name]
        log_path = out_dir / f'{name}.log'
        cmd = [
            sys.executable,
            'run.py',
            '--benchmark',
            str(one_question_path),
            '--checkpoint-every',
            '1',
        ]
        if runner.frame_budget_arg:
            cmd.extend(['--frame-budgets', str(args.frames)])
        cmd.extend(runner.extra_args)
        cmd.extend(runner_args.get(name, []))

        print(f'=== {name} ===')
        if args.dry_run:
            print(f'cwd: {runner.workdir}')
            print('cmd: ' + ' '.join(shlex.quote(x) for x in cmd))
            status = 'dry_run'
            code = 0
        else:
            code = _run_and_log(name, cmd, runner.workdir, env, log_path, quiet=args.quiet)
            status = 'ok' if code == 0 else f'error:{code}'

        result_row = _latest_result_row(output_root, runner.experiment, qid)
        pred = (result_row or {}).get('pred', '')
        gold = str(question.get('answer', '')).strip().upper()
        correct = bool(pred and pred.strip().upper() == gold)
        if code == 0 and not result_row and not args.dry_run:
            status = 'no_result'

        summary_row = {
            'runner': name,
            'status': status,
            'qid': qid,
            'video_file_id': question.get('video_file_id', ''),
            'model': (result_row or {}).get('model', ''),
            'backbone': (result_row or {}).get('backbone', ''),
            'num_frames': (result_row or {}).get('num_frames', ''),
            'gold': gold,
            'pred': pred,
            'correct': correct if pred else '',
            'latency_s': (result_row or {}).get('latency_s', ''),
            'results_csv': (result_row or {}).get('_results_csv', ''),
            'log': str(log_path),
        }
        summary.append(summary_row)
        print(f'{name}: status={status} pred={pred or "-"} correct={summary_row["correct"]}')
        print()

        if code != 0 and args.fail_fast:
            break

    _write_summary(summary, out_dir)
    _print_accuracy_summary(summary)
    print(f'Summary CSV: {out_dir / "summary.csv"}')
    print(f'Summary JSON: {out_dir / "summary.json"}')
    print(f'Accuracy JSON: {out_dir / "accuracy.json"}')

    all_ok = all(str(r['status']) in ('ok', 'dry_run') for r in summary)
    if args.fail_fast or args.require_all_ok:
        return 0 if all_ok else 1
    return 0


if __name__ == '__main__':
    raise SystemExit(main())