"""
Benchmark loader for CHNB.

Reads `data/benchmark.jsonl` (or any JSONL file with the same schema) into a
canonical pandas DataFrame. The public schema has nine fields per question:

    qid              str
    video_file_id    str
    category         str
    question         str
    option_A         str
    option_B         str
    option_C         str
    option_D         str
    answer           str   ('A' / 'B' / 'C' / 'D')
"""
from __future__ import annotations

import re
from pathlib import Path
from typing import Any

import pandas as pd


def normalize_answers(raw: Any) -> str:
    """Normalize an answer to a single uppercase letter ('A'/'B'/'C'/'D')."""
    if raw is None:
        return ''
    if isinstance(raw, (list, tuple, set)):
        items = [str(x) for x in raw]
    else:
        items = re.split(r'[,;\s/|]+', str(raw))
    for it in items:
        L = it.strip().upper()
        if L in {'A', 'B', 'C', 'D'}:
            return L
    return ''


def load_benchmark(path: str | Path) -> pd.DataFrame:
    """
    Load `data/benchmark.jsonl` (or any compatible JSONL) and return a canonical
    DataFrame. The returned DataFrame has the public schema's nine fields plus
    the aliases `video_id`, `gold`, and `q_type` for back-compat with the rest
    of the codebase.

    Args:
        path: path to a JSONL file with one question per line.

    Returns:
        DataFrame indexed 0..n-1 with string-typed columns. `gold` is a
        single uppercase letter; rows whose `answer` field doesn't normalize
        to one of {A, B, C, D} are dropped.
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f'Benchmark not found: {p}')
    df = pd.read_json(p, lines=True)

    required = ['qid', 'video_file_id', 'question',
                'option_A', 'option_B', 'option_C', 'option_D', 'answer']
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise ValueError(f'Missing required columns in {p.name}: {missing}')

    for col in ('qid', 'video_file_id', 'category', 'question',
                'option_A', 'option_B', 'option_C', 'option_D'):
        if col not in df.columns:
            df[col] = ''
        df[col] = df[col].fillna('').astype(str).str.strip()

    df['video_id'] = df['video_file_id']
    df['gold'] = df['answer'].apply(normalize_answers)
    df['q_type'] = 'SC'

    df = df[df['gold'].astype(bool)].reset_index(drop=True)

    return df


def unique_movies(df: pd.DataFrame) -> list[str]:
    """Unique video_file_ids in the order they first appear."""
    seen, out = set(), []
    for v in df['video_file_id']:
        if v and v not in seen:
            seen.add(v)
            out.append(v)
    return out