"""
BOBSL (BBC-Oxford British Sign Language) dataset loader.

BOBSL is a large-scale BSL dataset from BBC broadcasts with subtitles
and automatically spotted mouthings annotations.

Dataset structure:
    {SLTK_DATASETS_ROOT}/mixedmode/BOBSL/bobsl-v1-release/
        videos/*.mp4                          # Video files
        subtitles/audio-aligned/*.vtt         # WebVTT subtitle files
        spottings/mouthings.json              # Mouthings with train/val/test splits
        metadata_public_episodes.tsv          # Episode metadata

Mouthings format (word-centric):
    {
        "train": {
            "word": {
                "names": [video_id1, video_id2, ...],      # Video IDs
                "global_times": [timestamp1, timestamp2, ...],  # Center times
                "probs": [prob1, prob2, ...]              # Confidence scores
            },
            ...
        },
        ...
    }

Usage:
    from sltk.data.datasets import BOBSLDataset

    # Word-based dataset (each sample = one word's occurrences)
    # Uses default paths from environment or hardcoded defaults
    dataset = BOBSLDataset(
        split="train",
        annotation_type="mouthings",
        index_by="word",  # Default
    )

    # Video-based dataset (each sample = one video's mouthings)
    dataset = BOBSLDataset(
        split="train",
        annotation_type="mouthings",
        index_by="video",
    )

    # Subtitles (video-based)
    dataset = BOBSLDataset(
        split="train",
        annotation_type="subtitles",
    )
"""

import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Any

from sltk.data.core import Sample, Segment, SegmentList
from sltk.data.datasets.base import BaseDataset

# Pre-compiled regex patterns for better performance
# VTT timestamp format: 00:00:00.000 --> 00:00:00.000
_VTT_TIME_PATTERN = re.compile(r"(\d{2}:\d{2}:\d{2}\.\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}\.\d{3})")


def _get_default_video_root() -> Path:
    """Get default video root from config."""
    from sltk.config import get_default_dataset_roots

    return get_default_dataset_roots()["mixedmode"] / "BOBSL" / "bobsl-v1-release"


class BOBSLDataset(BaseDataset):
    """
    BOBSL dataset loader.

    Supports two annotation types:
    - mouthings: Automatically spotted word-level mouthings (has splits)
    - subtitles: Audio-aligned VTT subtitles

    For mouthings, supports two indexing modes:
    - word: Each sample is a word with all its occurrences across videos
    - video: Each sample is a video with all mouthings found in it
    """

    NAME = "bobsl"
    TASKS = ["mouthing_recognition", "subtitle_alignment", "sign_spotting"]
    LANGUAGES = ["bsl"]
    CITATION = "Albanie et al. (2021) BBC-Oxford British Sign Language Dataset"

    # Computed dynamically to respect environment variables
    @property
    def DEFAULT_VIDEO_ROOT(self) -> Path:  # type: ignore[override]
        return _get_default_video_root()

    # Estimated mouthing duration in seconds
    MOUTHING_DURATION = 0.5

    def __init__(
        self,
        root: str | Path | None = None,
        split: str = "train",
        feature_root: str | Path | None = None,
        load_poses: bool = False,
        pose_format: str = "mediapipe",
        fps: float = 25.0,
        annotation_type: str = "mouthings",
        index_by: str = "word",
        min_probability: float = 0.0,
    ):
        """
        Initialize BOBSL dataset.

        Args:
            root: Path to BOBSL directory
            split: Data split ('train', 'val', 'public_test')
            feature_root: Path to features directory (if any)
            load_poses: Whether to load pose features
            pose_format: Pose format to load
            fps: Video FPS
            annotation_type: 'mouthings' or 'subtitles'
            index_by: For mouthings - 'word' (each sample = word) or 'video' (each sample = video)
            min_probability: Minimum probability for mouthing annotations
        """
        root = root or _get_default_video_root()
        self.annotation_type = annotation_type
        self.index_by = index_by
        self.min_probability = min_probability

        # Load annotations before parent init (needed for index)
        self._root = Path(root)
        self._raw_annotations = self._load_raw_annotations(split)
        self._video_index = self._build_video_index()

        super().__init__(
            root=root,
            split=split,
            feature_root=feature_root,
            load_poses=load_poses,
            pose_format=pose_format,
            fps=fps,
        )

    def _load_raw_annotations(self, split: str) -> dict[str, Any]:
        """Load raw mouthings annotations from JSON."""
        if self.annotation_type == "mouthings":
            mouthings_path = self._root / "spottings" / "mouthings.json"
            if not mouthings_path.exists():
                return {}

            with open(mouthings_path) as f:
                data = json.load(f)

            # Map split names
            split_map = {
                "train": "train",
                "val": "val",
                "test": "public_test",
                "public_test": "public_test",
            }
            split_key = split_map.get(split, split)

            return data.get(split_key, {})

        return {}

    def _build_video_index(self) -> dict[str, list[dict]]:
        """Build video-based index from word-based mouthings."""
        video_index = defaultdict(list)

        for word, ann in self._raw_annotations.items():
            names = ann.get("names", [])  # Video IDs
            times = ann.get("global_times", [])  # Center timestamps
            probs = ann.get("probs", [])

            for i, video_id in enumerate(names):
                if i < len(times):
                    center_time = times[i]
                    prob = probs[i] if i < len(probs) else 1.0

                    if prob >= self.min_probability:
                        video_index[video_id].append(
                            {
                                "word": word,
                                "center_time": center_time,
                                "probability": prob,
                            }
                        )

        return dict(video_index)

    def _load_index(self) -> list[str]:
        """Build index of samples."""
        if self.annotation_type == "mouthings":
            if self.index_by == "word":
                # Filter to words with occurrences meeting probability threshold
                words = []
                for word, ann in self._raw_annotations.items():
                    probs = ann.get("probs", [])
                    if any(p >= self.min_probability for p in probs):
                        words.append(word)
                return sorted(words)
            else:  # index_by == "video"
                return sorted(self._video_index.keys())

        elif self.annotation_type == "subtitles":
            # All videos with VTT files
            vtt_dir = self.root / "subtitles" / "audio-aligned"
            if vtt_dir.exists():
                return [p.stem for p in sorted(vtt_dir.glob("*.vtt"))]
            return []

        return []

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single BOBSL sample."""
        if self.annotation_type == "mouthings":
            if self.index_by == "word":
                return self._load_word_sample(sample_id)
            else:
                return self._load_video_sample(sample_id)
        else:
            return self._load_subtitle_sample(sample_id)

    def _load_word_sample(self, word: str) -> Sample:
        """Load a word-based sample (all occurrences of a word)."""
        if word not in self._raw_annotations:
            return Sample(id=word, dataset=self.NAME, split=self.split)

        ann = self._raw_annotations[word]
        names = ann.get("names", [])
        times = ann.get("global_times", [])
        probs = ann.get("probs", [])

        segments = []
        for i, video_id in enumerate(names):
            if i < len(times):
                center_time = times[i]
                prob = probs[i] if i < len(probs) else 1.0

                if prob >= self.min_probability:
                    # Estimate segment boundaries from center time
                    start = max(0, center_time - self.MOUTHING_DURATION / 2)
                    end = center_time + self.MOUTHING_DURATION / 2

                    segments.append(
                        Segment(
                            start=start,
                            end=end,
                            label=word,
                            confidence=prob,
                            tier="mouthing",
                            metadata={"video_id": video_id},
                        )
                    )

        return Sample(
            id=word,
            video_path=None,  # Word sample spans multiple videos
            poses=None,
            segments=SegmentList(segments),
            glosses=[word],
            dataset=self.NAME,
            split=self.split,
            metadata={
                "num_occurrences": len(segments),
                "videos": list(set(names)),
            },
        )

    def _load_video_sample(self, video_id: str) -> Sample:
        """Load a video-based sample (all mouthings in a video)."""
        # Video path
        video_path_candidate = self.root / "videos" / f"{video_id}.mp4"
        video_path: Path | None = video_path_candidate if video_path_candidate.exists() else None

        # Get mouthings for this video
        mouthings = self._video_index.get(video_id, [])

        segments = []
        words = set()
        for m in mouthings:
            word = m["word"]
            center_time = m["center_time"]
            prob = m["probability"]

            start = max(0, center_time - self.MOUTHING_DURATION / 2)
            end = center_time + self.MOUTHING_DURATION / 2

            segments.append(
                Segment(
                    start=start,
                    end=end,
                    label=word,
                    confidence=prob,
                    tier="mouthing",
                )
            )
            words.add(word)

        # Sort by time
        segments.sort(key=lambda s: s.start)

        # Load metadata
        metadata = self._load_episode_metadata(video_id)

        return Sample(
            id=video_id,
            video_path=video_path,
            poses=None,
            segments=SegmentList(segments),
            glosses=sorted(words),
            text=metadata.get("synopsis"),
            dataset=self.NAME,
            split=self.split,
            metadata=metadata,
        )

    def _load_subtitle_sample(self, video_id: str) -> Sample:
        """Load a subtitle-based sample."""
        video_path_candidate = self.root / "videos" / f"{video_id}.mp4"
        video_path: Path | None = video_path_candidate if video_path_candidate.exists() else None

        segments = self._load_subtitles(video_id)
        glosses = list(set(s.label for s in segments if s.label))
        metadata = self._load_episode_metadata(video_id)

        return Sample(
            id=video_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=glosses,
            text=metadata.get("synopsis"),
            dataset=self.NAME,
            split=self.split,
            metadata=metadata,
        )

    def _load_subtitles(self, video_id: str) -> SegmentList:
        """Load VTT subtitles for a video."""
        vtt_path = self.root / "subtitles" / "audio-aligned" / f"{video_id}.vtt"
        if not vtt_path.exists():
            return SegmentList()

        segments = []

        with open(vtt_path, encoding="utf-8") as f:
            content = f.read()

        # Parse VTT format using pre-compiled pattern
        lines = content.split("\n")
        i = 0
        while i < len(lines):
            line = lines[i].strip()
            match = _VTT_TIME_PATTERN.match(line)

            if match:
                start = self._parse_vtt_time(match.group(1))
                end = self._parse_vtt_time(match.group(2))

                # Get subtitle text (following lines until empty)
                text_lines = []
                i += 1
                while i < len(lines) and lines[i].strip():
                    text_lines.append(lines[i].strip())
                    i += 1

                text = " ".join(text_lines)
                if text:
                    segments.append(
                        Segment(
                            start=start,
                            end=end,
                            label=text,
                            tier="subtitle",
                        )
                    )
            i += 1

        return SegmentList(segments)

    def _parse_vtt_time(self, time_str: str) -> float:
        """Parse VTT timestamp to seconds."""
        parts = time_str.split(":")
        hours = int(parts[0])
        minutes = int(parts[1])
        seconds = float(parts[2])
        return hours * 3600 + minutes * 60 + seconds

    def _load_episode_metadata(self, video_id: str) -> dict[str, Any]:
        """Load episode metadata from TSV."""
        metadata_path = self.root / "metadata_public_episodes.tsv"
        if not metadata_path.exists():
            return {}

        try:
            with open(metadata_path, encoding="utf-8") as f:
                header = f.readline().strip().split("\t")
                for line in f:
                    parts = line.strip().split("\t")
                    if len(parts) > 0:
                        row = dict(zip(header, parts))
                        if video_id in str(row):
                            return row
        except Exception:
            pass

        return {}

    def get_vocabulary(self, min_count: int = 1) -> list[str]:
        """Get unique words from mouthings."""
        if self.annotation_type != "mouthings":
            return []

        word_counts: dict[str, int] = {}
        for word, ann in self._raw_annotations.items():
            probs = ann.get("probs", [])
            count = sum(1 for p in probs if p >= self.min_probability)
            if count > 0:
                word_counts[word] = count

        return sorted(w for w, c in word_counts.items() if c >= min_count)

    def get_split_stats(self) -> dict[str, Any]:
        """Get statistics for current split."""
        if self.annotation_type != "mouthings":
            return {"num_samples": len(self._index)}

        total_occurrences = 0
        unique_videos = set()

        for word, ann in self._raw_annotations.items():
            names = ann.get("names", [])
            probs = ann.get("probs", [])

            for i, video_id in enumerate(names):
                prob = probs[i] if i < len(probs) else 1.0
                if prob >= self.min_probability:
                    total_occurrences += 1
                    unique_videos.add(video_id)

        return {
            "num_words": len(self.get_vocabulary()),
            "num_videos": len(unique_videos),
            "total_occurrences": total_occurrences,
            "index_by": self.index_by,
            "num_samples": len(self._index),
        }
