"""
Path configuration for CHNB.

Layout:
    chnb/
    ├── data/
    │   ├── benchmark.jsonl
    │   ├── non_benchmark.jsonl
    │   ├── movie_metadata.tsv
    │   ├── mp4/                <- populated by scripts/download_movies.sh
    │   ├── frames/             <- pre-computed frame extractions (Socratic / VideoAgent)
    │   ├── clip_embeddings/    <- CLIP features (VideoAgent)
    │   ├── clips/              <- shot-grouped video chunks (Socratic clip captions)
    │   ├── captions/           <- per-chunk captions (Socratic)
    │   └── subtitles/          <- per-movie ASR / OpenSubtitles transcripts
    │       (one .txt per video_file_id)
    ├── output/
    │   ├── results/            <- one CSV per (experiment, model, config)
    │   └── world_state_histories/  <- cached chronological caption assemblies (Socratic)
    └── ...

All paths are overridable via env vars:
    CHNB_PROJECT_ROOT          — override the project root
    CHNB_DATA_ROOT             — override data/   (default: $PROJECT_ROOT/data)
    CHNB_OUTPUT_ROOT           — override output/ (default: $PROJECT_ROOT/output)
    CHNB_BENCHMARK             — override benchmark JSONL
    CHNB_MOVIE_METADATA        — override movie metadata TSV
    CHNB_VIDEO_DIR             — source mp4s
    CHNB_SUBTITLES_DIR         — subtitles dir
    CHNB_FRAMES_DIR            — pre-extracted frames
    CHNB_CLIP_EMBEDDINGS_DIR   — CLIP feature cache
    CHNB_CLIPS_DIR             — extracted clips
    CHNB_CAPTIONS_DIR          — generated captions
    CHNB_SHOTS_DIR             — shot-boundary files
    CHNB_RESULTS_DIR           — result CSVs
    CHNB_WSH_DIR               — world-state history cache

Tilde-expansion is applied automatically.
"""
from __future__ import annotations

import os
from pathlib import Path

from util.config import (
    BENCHMARK_ENV,
    CAPTIONS_DIR_ENV,
    CLIP_EMBEDDINGS_DIR_ENV,
    CLIPS_DIR_ENV,
    DATA_ROOT_ENV,
    FRAMES_DIR_ENV,
    GCS_BUCKET_ENV,
    MOVIE_METADATA_ENV,
    OUTPUT_ROOT_ENV,
    PROJECT_ROOT_ENV,
    RESULTS_DIR_ENV,
    SHOTS_DIR_ENV,
    SUBTITLES_DIR_ENV,
    VIDEO_DIR_ENV,
    WSH_DIR_ENV,
    load_public_env,
)


def _p(env: str, default: str | Path) -> Path:
    """Env var or default, with ~-expansion."""
    return Path(os.environ.get(env, str(default))).expanduser()


load_public_env(Path(__file__).resolve().parents[1])

PROJECT_ROOT = _p(PROJECT_ROOT_ENV, Path(__file__).resolve().parents[1])

DATA_ROOT = _p(DATA_ROOT_ENV, PROJECT_ROOT / 'data')
OUTPUT_ROOT = _p(OUTPUT_ROOT_ENV, PROJECT_ROOT / 'output')

BENCHMARK_PATH = _p(BENCHMARK_ENV, DATA_ROOT / 'benchmark.jsonl')
MOVIE_METADATA_PATH = _p(MOVIE_METADATA_ENV, DATA_ROOT / 'movie_metadata.tsv')

VIDEO_DIR = _p(VIDEO_DIR_ENV, DATA_ROOT / 'mp4')
SUBTITLES_DIR = _p(SUBTITLES_DIR_ENV, DATA_ROOT / 'subtitles')

FRAMES_DIR = _p(FRAMES_DIR_ENV, DATA_ROOT / 'frames')
CLIP_DIR = _p(CLIP_EMBEDDINGS_DIR_ENV, DATA_ROOT / 'clip_embeddings')
CLIPS_DIR = _p(CLIPS_DIR_ENV, DATA_ROOT / 'clips')
CAPTIONS_DIR = _p(CAPTIONS_DIR_ENV, DATA_ROOT / 'captions')
SHOTS_DIR = _p(SHOTS_DIR_ENV, DATA_ROOT / 'shots')
RESULTS_DIR = _p(RESULTS_DIR_ENV, OUTPUT_ROOT / 'results')
WSH_DIR = _p(WSH_DIR_ENV, OUTPUT_ROOT / 'world_state_histories')

GCS_BUCKET = os.environ.get(GCS_BUCKET_ENV, '')


def get_frames_dir(video_file_id: str, kind: str, n: int | float) -> Path:
    """
    kind='fps'     -> frames/{vid}/fps_{n}/     (e.g. fps_0.5)
    kind='uniform' -> frames/{vid}/uniform_{n}/ (e.g. uniform_128)
    """
    return FRAMES_DIR / video_file_id / f'{kind}_{n}'


def get_captions_dir(captioner_dir_name: str, video_file_id: str) -> Path:
    """captions/{captioner_dir_name}/{vid}/"""
    return CAPTIONS_DIR / captioner_dir_name / video_file_id


def get_clips_dir(video_file_id: str, chunk_mode: str) -> Path:
    """clips/{chunk_mode}/{vid}/  — each video's extracted clips + manifest.json"""
    return CLIPS_DIR / chunk_mode / video_file_id


def get_clip_path(video_file_id: str, kind: str, n: int | float) -> Path:
    """clip_embeddings/{vid}_{kind}_{n}.npy  — CLIP feature cache."""
    return CLIP_DIR / f'{video_file_id}_{kind}_{n}.npy'


def get_wsh_path(variant: str, video_file_id: str) -> Path:
    """output/world_state_histories/{variant}/{vid}.txt"""
    return WSH_DIR / variant / f'{video_file_id}.txt'


def model_slug(model_id: str) -> str:
    """
    Filesystem-safe slug for a model id, used in result filenames.

        'gpt-5-mini'                     -> 'gpt-5-mini'
        'Qwen/Qwen3-VL-8B-Instruct'      -> 'Qwen-Qwen3-VL-8B-Instruct'
        'THUDM/GLM-4.1V-9B-Thinking'     -> 'THUDM-GLM-4.1V-9B-Thinking'

    Anything that isn't [A-Za-z0-9._-] is collapsed to a single dash; leading
    and trailing dashes are stripped.
    """
    import re
    if not model_id:
        return ''
    s = re.sub(r'[^A-Za-z0-9._-]+', '-', str(model_id))
    return s.strip('-')


def results_filename(model_id: str = '', *, config: str = '',
                     backbone: str = '') -> str:
    """
    Canonical results-CSV filename:

        results_{backbone}_{model-slug}__{config}.csv

    Empty segments are omitted. The double underscore between model and config
    keeps `{model}__{config}` parseable when auditing later.

        results_filename(model_id='gpt-5-mini')
            -> 'results_gpt-5-mini.csv'
        results_filename(backbone='gpt', model_id='gpt-5-mini')
            -> 'results_gpt_gpt-5-mini.csv'
        results_filename(model_id='gemini-3-flash-preview', config='frames_128')
            -> 'results_gemini-3-flash-preview__frames_128.csv'
    """
    parts = ['results']
    if backbone:
        parts.append(backbone)
    slug = model_slug(model_id)
    if slug:
        parts.append(slug)
    name = '_'.join(parts)
    if config:
        sep = '__' if slug else '_'
        name = f'{name}{sep}{config}'
    return f'{name}.csv'


def results_filename_for_budget(num_frames: int, model_id: str = '') -> str:
    """
    Per-budget CSV filename. Frame-budget-driven runners use this so each
    budget gets its own CSV (text_only / frames_64 / frames_128 / ...).

        results_filename_for_budget(0,   'gpt-5-mini')   -> 'results_gpt-5-mini__text_only.csv'
        results_filename_for_budget(128, 'gpt-5-mini')   -> 'results_gpt-5-mini__frames_128.csv'
    """
    config = 'text_only' if num_frames == 0 else f'frames_{num_frames}'
    return results_filename(model_id=model_id, config=config)


def ensure_dir(path: Path) -> Path:
    """mkdir -p then return the path."""
    path.mkdir(parents=True, exist_ok=True)
    return path


def summary() -> str:
    """Human-readable dump of all resolved paths; useful in `python -m util.paths`."""
    rows = [
        ('PROJECT_ROOT',    PROJECT_ROOT),
        ('DATA_ROOT',       DATA_ROOT),
        ('OUTPUT_ROOT',     OUTPUT_ROOT),
        ('BENCHMARK_PATH',  BENCHMARK_PATH),
        ('MOVIE_METADATA',  MOVIE_METADATA_PATH),
        ('VIDEO_DIR',       VIDEO_DIR),
        ('SUBTITLES_DIR',   SUBTITLES_DIR),
        ('FRAMES_DIR',      FRAMES_DIR),
        ('CLIP_DIR',        CLIP_DIR),
        ('CLIPS_DIR',       CLIPS_DIR),
        ('CAPTIONS_DIR',    CAPTIONS_DIR),
        ('RESULTS_DIR',     RESULTS_DIR),
        ('WSH_DIR',         WSH_DIR),
    ]
    return '\n'.join(f'{k:18s} = {v}' for k, v in rows)


if __name__ == '__main__':
    print(summary())