"""
Subtitle loader for the subtitles_only paradigm.

Supports two on-disk formats at `{SUBTITLES_DIR}/{video_file_id}.speech.txt`:

       LANGS\\t[["en", "0.703"], ["en", "0.712"], ...]

   one (lang, confidence) tuple per following segment. Each subsequent
   line is tab-separated:

       col 0: track / speaker id          (e.g. '1')
       col 1: segment metric              (e.g. '12' — word count or similar)
       col 2: start_sec                   (float, e.g. '4.37713')
       col 3: end_sec                     (float, e.g. '4.91713')
       col 4: (empty)
       col 5: text                        (e.g. 'Oh, my God.')
       col 6: word-level alignment JSON   ([{"w": ..., "s": ..., "e": ...}, ...])

We follow MF2's `extract_srt_text()` treatment exactly: keep dialog text,
drop metadata, space-join in segment order. For both formats that means
take the text column (col 5 / `text`) verbatim. We do NOT filter on
confidence, language, et.c as in MF2.

Format detection: if the first line starts with `LANGS\\t`, parse as
iaspeech; otherwise fall back to Whisper TSV via pandas.

Cf. https://github.com/deep-spin/MF2
"""
from __future__ import annotations

from io import StringIO
from pathlib import Path
from typing import Iterable

import pandas as pd

from util.paths import SUBTITLES_DIR


def subtitles_path(video_file_id: str, root: Path | None = None) -> Path:
    """Resolve the on-disk path for a video's transcript."""
    return (root or SUBTITLES_DIR) / f'{video_file_id}.speech.txt'


# Column index of the dialog text in the iaspeech 7-column TSV.
_IASPEECH_TEXT_COL = 5


def _parse_iaspeech_text(raw: str) -> str:
    """Parse the iaspeech format: drop the LANGS header, join col-5 text
    in segment order. Skips empty / malformed rows silently."""
    lines = raw.splitlines()
    if not lines:
        return ''
    # Drop the LANGS metadata line if present (always first).
    if lines[0].lstrip().startswith('LANGS\t'):
        lines = lines[1:]
    out: list[str] = []
    for line in lines:
        if not line.strip():
            continue
        parts = line.split('\t')
        if len(parts) <= _IASPEECH_TEXT_COL:
            continue
        seg = parts[_IASPEECH_TEXT_COL].strip()
        if seg:
            out.append(seg)
    return ' '.join(out)


def _parse_whisper_tsv(raw: str) -> str:
    try:
        df = pd.read_csv(StringIO(raw), sep='\t')
    except (pd.errors.EmptyDataError, pd.errors.ParserError):
        return ''
    if 'text' not in df.columns or len(df) == 0:
        return ''
    texts = df['text'].fillna('').astype(str).str.strip()
    texts = texts[texts != '']
    return ' '.join(texts.tolist())


def load_subtitles(video_file_id: str, *, root: Path | None = None) -> str:
    """Read the transcript for `video_file_id` and return the flat dialog
    string (space-joined per-segment text, MF2-style).

    Returns '' if the file is missing or unparseable. The caller decides
    what to do with empty subtitles (typical: skip the question or let
    the model commit to an answer with no context).
    """
    path = subtitles_path(video_file_id, root=root)
    if not path.exists():
        return ''
    try:
        raw = path.read_text(encoding='utf-8', errors='ignore')
    except OSError:
        return ''
    if not raw.strip():
        return ''
    # Format-detect on the first non-empty line.
    first_line = raw.lstrip().split('\n', 1)[0]
    if first_line.startswith('LANGS\t'):
        return _parse_iaspeech_text(raw)
    return _parse_whisper_tsv(raw)


def list_movies_with_subtitles(root: Path | None = None) -> list[str]:
    """Return the video_file_ids that have a `.speech.txt` on disk."""
    root = root or SUBTITLES_DIR
    if not root.exists():
        return []
    return sorted(p.stem.removesuffix('.speech') if p.name.endswith('.speech.txt')
                  else p.stem
                  for p in root.glob('*.speech.txt'))


def subtitles_coverage(vids: Iterable[str], root: Path | None = None
                       ) -> tuple[int, int]:
    """(n_with_subs, n_total) for the given list of video_file_ids."""
    vids = list(vids)
    have = sum(1 for v in vids if subtitles_path(v, root=root).exists())
    return have, len(vids)