"""
World-state history (WSH) builder for the Socratic paradigm.

Cf. https://github.com/keshik6/HourVideo

Usage:

    from util.world_state import build_world_state_history

    wsh = build_world_state_history('gemini_frames', 'ia.FooMovie_tt123')
    # wsh is a long string of
    #   **Timestamp**: HH:MM:SS - HH:MM:SS
    #   **Caption**: ...
    # blocks, separated by blank lines.
"""
from __future__ import annotations

import json
from pathlib import Path

from util.paths import CAPTIONS_DIR, WSH_DIR, get_captions_dir, get_clips_dir


def _seconds_to_hms(sec: float) -> str:
    h = int(sec // 3600)
    m = int((sec % 3600) // 60)
    s = int(sec % 60)
    return f'{h:02d}:{m:02d}:{s:02d}'


def _chunk_mode_from_variant(variant: str) -> str:
    """Map a captioner variant name to its clip chunk_mode, for manifest lookup."""
    if variant.endswith('_shot_grouped'):
        return 'shot_grouped'
    if variant.endswith('_fixed'):
        return 'fixed'
    return 'frames_1min'   # sentinel — no clip manifest; use 1-min fixed intervals


def build_world_state_history(
    variant: str,
    video_file_id: str,
    *,
    use_cache: bool = True,
) -> str:
    """
    Build a chronological world-state history string from pre-computed captions.

    Args:
        variant: captioner directory name under `data/captions/` (e.g.
            `gemini_frames`, `gemini_video_clip_shot_grouped`).
        video_file_id: movie identifier (matches the mp4 filename stem).
        use_cache: if True (default), read/write
            `output/world_state_histories/{variant}/{vid}.txt`.

    Returns:
        Multi-line string with per-chunk timestamp + caption, or '' if no
        captions exist for this (variant, video) pair.
    """
    cache_path = WSH_DIR / variant / f'{video_file_id}.txt'
    if use_cache and cache_path.exists():
        return cache_path.read_text()

    cap_dir = get_captions_dir(variant, video_file_id)
    if not cap_dir.exists():
        return ''

    txt_files = sorted(
        [p for p in cap_dir.glob('*.txt') if p.stem.isdigit()],
        key=lambda f: int(f.stem),
    )
    if not txt_files:
        return ''

    # Timestamps come from the clip manifest for shot-grouped / fixed mp4
    # variants, else fall back to 1-min synthetic intervals.
    timestamps: dict[int, tuple[float, float]] | None = None
    chunk_mode = _chunk_mode_from_variant(variant)
    if chunk_mode in ('shot_grouped', 'fixed'):
        clips_dir = get_clips_dir(video_file_id, chunk_mode)
        manifest_path = clips_dir / 'manifest.json'
        if manifest_path.exists():
            m = json.loads(manifest_path.read_text())
            timestamps = {i: (item['start'], item['end']) for i, item in enumerate(m)}

    segments = []
    for fname in txt_files:
        idx = int(fname.stem)
        caption = fname.read_text().strip()
        if timestamps and idx in timestamps:
            start_sec, end_sec = timestamps[idx]
        else:
            start_sec, end_sec = idx * 60.0, (idx + 1) * 60.0
        segments.append(
            f'**Timestamp**: {_seconds_to_hms(start_sec)} - {_seconds_to_hms(end_sec)}\n'
            f'**Caption**: {caption}\n'
        )
    wsh = '\n'.join(segments)

    if use_cache:
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        cache_path.write_text(wsh)
    return wsh


def list_captioners() -> list[str]:
    """Return every subdirectory name under `data/captions/`."""
    if not CAPTIONS_DIR.exists():
        return []
    return sorted(d.name for d in CAPTIONS_DIR.iterdir() if d.is_dir())


def captioner_coverage(variant: str) -> tuple[int, int]:
    """
    Return (n_movies_captioned, n_total_caption_files) for a variant.

    Useful for the socratic experiment's pre-flight check so you know which
    captioners are actually ready to run.
    """
    root = CAPTIONS_DIR / variant
    if not root.exists():
        return 0, 0
    movies = [d for d in root.iterdir() if d.is_dir()]
    total = sum(
        sum(1 for p in d.glob('*.txt') if p.stem.isdigit())
        for d in movies
    )
    return len(movies), total