#!/usr/bin/env python
"""
Socratic paradigm: per-question QA over a pre-computed world-state history.

Usage:
    # Default: every captioner present on disk
    python run.py

    # Specific combo:
    python run.py --captioners gemini_frames --qa-backbones gemini gpt claude
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

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

from util.backbones import call_backbone, init_backbone
from util.benchmark import load_benchmark
from util.config import (
    ANTHROPIC_API_KEY_ENV,
    DEFAULT_MODELS,
    GOOGLE_CLOUD_PROJECT_ENV,
    OPENAI_API_KEY_ENV,
)
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,
)
from util.prompts import HEADER_WORLD_STATE, mcq_prompt_json
from util.world_state import build_world_state_history, captioner_coverage, list_captioners


EXPERIMENT_NAME = 'socratic'
SUPPORTED_BACKBONES = ('gemini', 'gpt', 'claude', 'qwen3vl', 'qwen3', 'glm4v', 'glm4')
TEXT_ONLY_VARIANTS = ('qwen3', 'glm4')


def _init_backbone_guard(backbone: str, model_override: str | None = None) -> None:
    """Initialize a backbone; secrets are loaded from .env automatically by util.backbones."""
    from util.secrets import require
    if backbone == 'gpt':
        require(OPENAI_API_KEY_ENV)
    if backbone == 'claude':
        require(ANTHROPIC_API_KEY_ENV)

    kwargs = {}
    if model_override:
        kwargs['model_id'] = model_override
    if backbone == 'gemini':
        init_backbone('gemini',
                      project_id=os.environ.get(GOOGLE_CLOUD_PROJECT_ENV),
                      **kwargs)
    else:
        init_backbone(backbone, **kwargs)


def _print_captioner_inventory(requested: list[str]) -> list[str]:
    """Report which requested captioners are actually ready; return the subset that is."""
    available = set(list_captioners())
    ready: list[str] = []
    print('Captioner inventory:')
    for v in requested:
        n_movies, n_caps = captioner_coverage(v)
        marker = 'ok' if v in available and n_caps > 0 else 'no'
        print(f'  {marker} {v:40s} {n_movies} movies, {n_caps} caption files')
        if v in available and n_caps > 0:
            ready.append(v)
    if not ready:
        sys.exit('No captioners are ready. Run preprocess captioning first '
                 '(see experiments/preprocess/).')
    return ready


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--captioners', nargs='+', default=None,
                    help='Captioner dir names under data/captions/. '
                         'Default: every directory present on disk.')
    ap.add_argument('--qa-backbones', nargs='+',
                    default=['gemini'],
                    choices=SUPPORTED_BACKBONES,
                    help='One or more backbones to run QA with.')
    ap.add_argument('--include-text-only-variants', action='store_true')
    ap.add_argument('--model-override', nargs='+', metavar='backbone=model_id',
                    default=[],
                    help='Override default model per backbone.')
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--max-output-tokens', type=int, default=24576)
    ap.add_argument('--temperature', type=float, default=0.0)
    ap.add_argument('--checkpoint-every', type=int, default=50)
    args = ap.parse_args()

    if args.include_text_only_variants:
        for v in TEXT_ONLY_VARIANTS:
            if v not in args.qa_backbones:
                args.qa_backbones.append(v)
        print(f'  --include-text-only-variants: extended to '
              f'{" ".join(args.qa_backbones)}')

    requested = args.captioners if args.captioners else list_captioners()
    if not requested:
        sys.exit('No --captioners specified and data/captions/ is empty.')
    captioners = _print_captioner_inventory(requested)

    out_dir = ensure_dir(RESULTS_DIR / EXPERIMENT_NAME)
    df = load_benchmark(args.benchmark)
    print(f'\nBenchmark: {len(df)} questions across {df["video_file_id"].nunique()} movies')

    overrides = {}
    for kv in args.model_override:
        if '=' not in kv:
            sys.exit(f'--model-override entries must be backbone=model_id, got: {kv!r}')
        bb, mid = kv.split('=', 1)
        overrides[bb.strip()] = mid.strip()
    if 'claude' in args.qa_backbones and 'claude' not in overrides:
        overrides['claude'] = DEFAULT_MODELS['claude']
        print(f'  (auto-selecting {DEFAULT_MODELS["claude"]} for the claude backbone — '
              'Sonnet rate-limits eat Socratic\'s long prompts; pass '
              '--model-override claude=claude-sonnet-4-5 to force Sonnet.)')
    args._overrides = overrides

    # Backbone outer / captioner inner: lazy init means we never hold two
    # local models in GPU memory at once, and a single bad backbone (key,
    # retired model) doesn't block the rest of the sweep.
    for backbone in args.qa_backbones:
        print(f'--- Initializing backbone: {backbone} ---')
        try:
            _init_backbone_guard(backbone, overrides.get(backbone))
        except Exception as e:
            print(f'  init failed for {backbone}: {type(e).__name__}: {e!r}')
            print(f'  skipping {backbone}; other backbones will continue.')
            continue
        for captioner in captioners:
            _run_one(df, captioner, backbone, args, out_dir)


def _run_one(df, captioner: str, backbone: str, args, out_dir: Path) -> None:
    print(f'\n=== {EXPERIMENT_NAME}: captioner={captioner} x backbone={backbone} ===')

    def build_movie_ctx(vid, movie_df):
        wsh = build_world_state_history(captioner, vid)
        if not wsh:
            print(f'  SKIP {vid}: no captions under data/captions/{captioner}/{vid}/')
            return {}
        return {'wsh': wsh}

    def call_fn(row, ctx):
        wsh = ctx.get('wsh')
        if not wsh:
            return ''
        header = HEADER_WORLD_STATE.format(wsh=wsh)
        prompt = mcq_prompt_json(row, header=header)
        raw = call_backbone(
            backbone, prompt, images=None,
            system_instruction=('You answer movie MCQs based on a text description. '
                                'Answer ONLY from the information provided. Output valid JSON.'),
            json_mode=True,
            max_output_tokens=args.max_output_tokens,
            temperature=args.temperature,
        )
        return parse_mcq_json(raw, qid=row['qid'])

    model_id = _model_id_for(backbone)
    results_path = out_dir / f'results_{captioner}_{backbone}.csv'
    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        results_path=results_path,
        experiment_name=f'{EXPERIMENT_NAME} {captioner}x{backbone}',
        extra_row_fields={
            'caption_variant': captioner,
            'backbone':        backbone,
            'model':           model_id,
        },
        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)


if __name__ == '__main__':
    main()