"""
YouTube Sign Language (YTSL-25 and YT-ASL) dataset loaders.

These datasets contain sign language videos collected from YouTube
with caption-based annotations (not glosses).

YTSL-25 structure:
    /path/to/ytsl-25/
        videos/
            *.mp4
        captions/
            *.tsv  (video_id, start, end, caption)
        metadata.json

YT-ASL structure:
    /path/to/yt-asl/
        videos/
            *.mp4
        subtitles/
            *.vtt or *.srt
        manifest.json

Reference:
    Uthus, D., et al. "YouTube-SL-25: A Large-Scale, Open-Domain
    Multilingual Sign Language Parallel Corpus." 2023.

Usage:
    from sltk.data.datasets import YTSL25Dataset, YTASLDataset

    # YTSL-25 (multilingual)
    dataset = YTSL25Dataset(
        root="/path/to/ytsl-25",
        language="asl",  # or 'bsl', 'dgs', etc.
    )

    # YT-ASL
    dataset = YTASLDataset(
        root="/path/to/yt-asl",
    )

    sample = dataset[0]
    print(sample.translation)  # Caption text
"""

import csv
import json
import re
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
_TIME_HMS_PATTERN = re.compile(r"(\d+):(\d+):(\d+)(?:\.(\d+))?")  # HH:MM:SS.mmm
_TIME_MS_PATTERN = re.compile(r"(\d+):(\d+)(?:\.(\d+))?")  # MM:SS.mmm
_TIME_CHECK_PATTERN = re.compile(r"\d+:\d+")  # Quick check for time line
_VTT_TIME_PATTERN = re.compile(
    r"(\d+:\d+:\d+\.\d+|\d+:\d+\.\d+)\s*-->\s*(\d+:\d+:\d+\.\d+|\d+:\d+\.\d+)"
)
_SRT_TIME_PATTERN = re.compile(r"(\d+:\d+:\d+[,\.]\d+)\s*-->\s*(\d+:\d+:\d+[,\.]\d+)")


class YTSL25Dataset(BaseDataset):
    """
    YouTube Sign Language 25 (YTSL-25) dataset loader.

    Multilingual sign language dataset with 25 languages.
    Uses caption-based annotations (not glosses).
    """

    NAME = "ytsl25"
    TASKS = ["sign_language_translation", "caption_based_recognition"]
    LANGUAGES = [
        "asl",
        "bsl",
        "dgs",
        "lsf",
        "lis",
        "libras",
        "jsl",
        "ksl",
        "auslan",
        "nzsl",
        "csl",
        "tsl",
        "hksl",
        "isl",
        "psl",
        "rsl",
        "ukr_sl",
        "asl_arab",
        "gsl",
        "ssl",
        "fsl",
        "lse",
        "ngsl",
        "nsl",
        "vgt",
    ]
    CITATION = """
    @article{uthus2023youtube,
        title={YouTube-SL-25: A Large-Scale, Open-Domain Multilingual Sign Language Parallel Corpus},
        author={Uthus, David and others},
        year={2023}
    }
    """

    def __init__(
        self,
        root: str | Path,
        split: str = "train",
        feature_root: str | Path | None = None,
        load_poses: bool = False,
        pose_format: str = "mediapipe",
        fps: float = 25.0,
        language: str = "asl",
    ):
        """
        Initialize YTSL-25 dataset.

        Args:
            root: Path to YTSL-25 directory
            split: Data split ('train', 'val', 'test')
            feature_root: Path to pre-extracted features
            load_poses: Whether to load pose data
            pose_format: Pose format to load
            fps: Video FPS
            language: Sign language to load ('asl', 'bsl', etc.)
        """
        self.language = language.lower()

        self._root = Path(root)
        self._captions = self._load_captions()
        self._metadata = self._load_metadata()

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

    def _load_captions(self) -> dict[str, list[dict[str, Any]]]:
        """Load caption TSV files."""
        captions = {}

        caption_dir = self._root / "captions"
        if not caption_dir.exists():
            caption_dir = self._root / self.language / "captions"

        if not caption_dir.exists():
            # Try loading single caption file
            for pattern in ["captions.tsv", "*.tsv"]:
                for tsv_file in self._root.glob(pattern):
                    captions.update(self._parse_tsv(tsv_file))
            return captions

        # Load all TSV files in caption directory
        for tsv_file in caption_dir.glob("*.tsv"):
            captions.update(self._parse_tsv(tsv_file))

        return captions

    def _parse_tsv(self, tsv_path: Path) -> dict[str, list[dict[str, Any]]]:
        """Parse a caption TSV file."""
        captions: dict[str, list[dict[str, Any]]] = {}

        with open(tsv_path, encoding="utf-8", errors="replace") as f:
            reader = csv.reader(f, delimiter="\t")
            header = next(reader, None)

            if header is None:
                return captions

            # Determine column indices
            col_map = {}
            for i, col in enumerate(header):
                col_lower = col.lower().strip()
                if "video" in col_lower or "id" in col_lower:
                    col_map["video_id"] = i
                elif "start" in col_lower:
                    col_map["start"] = i
                elif "end" in col_lower:
                    col_map["end"] = i
                elif "text" in col_lower or "caption" in col_lower or "subtitle" in col_lower:
                    col_map["text"] = i
                elif "lang" in col_lower:
                    col_map["language"] = i

            for row in reader:
                if len(row) < len(col_map):
                    continue

                video_id = row[col_map.get("video_id", 0)] if "video_id" in col_map else ""
                if not video_id:
                    continue

                # Filter by language if specified
                if "language" in col_map:
                    lang = row[col_map["language"]].lower()
                    if lang != self.language and self.language != "all":
                        continue

                start = (
                    self._parse_time(row[col_map.get("start", 1)]) if "start" in col_map else 0.0
                )
                end = self._parse_time(row[col_map.get("end", 2)]) if "end" in col_map else 0.0
                text = row[col_map.get("text", 3)] if "text" in col_map else ""

                if video_id not in captions:
                    captions[video_id] = []

                captions[video_id].append(
                    {
                        "start": start,
                        "end": end,
                        "text": text,
                    }
                )

        return captions

    def _parse_time(self, time_str: str) -> float:
        """Parse time string to seconds."""
        if not time_str:
            return 0.0

        try:
            return float(time_str)
        except ValueError:
            pass

        # Try HH:MM:SS.mmm format
        match = _TIME_HMS_PATTERN.match(time_str)
        if match:
            h, m, s, ms = match.groups()
            return int(h) * 3600 + int(m) * 60 + int(s) + (int(ms or 0) / 1000)

        # Try MM:SS.mmm format
        match = _TIME_MS_PATTERN.match(time_str)
        if match:
            m, s, ms = match.groups()
            return int(m) * 60 + int(s) + (int(ms or 0) / 1000)

        return 0.0

    def _load_metadata(self) -> dict[str, dict[str, Any]]:
        """Load metadata JSON."""
        metadata = {}

        meta_paths = [
            self._root / "metadata.json",
            self._root / self.language / "metadata.json",
            self._root / "manifest.json",
        ]

        for meta_path in meta_paths:
            if meta_path.exists():
                with open(meta_path, encoding="utf-8") as f:
                    data = json.load(f)
                    if isinstance(data, list):
                        for item in data:
                            vid_id = item.get("video_id", "") or item.get("id", "")
                            metadata[vid_id] = item
                    else:
                        metadata = data
                break

        return metadata

    def _load_index(self) -> list[str]:
        """Get sample IDs (video IDs with captions)."""
        # Filter by split if metadata has split info
        video_ids = []
        for vid_id, caps in self._captions.items():
            if not caps:
                continue

            meta = self._metadata.get(vid_id, {})
            vid_split = meta.get("split", "train")

            if vid_split == self.split or self.split == "all":
                video_ids.append(vid_id)

        return video_ids

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single YTSL-25 sample."""
        captions = self._captions.get(sample_id, [])
        meta = self._metadata.get(sample_id, {})

        # Find video
        video_path = self._find_video(sample_id)

        # Create segments from captions
        segments = SegmentList(
            [
                Segment(
                    start=cap["start"],
                    end=cap["end"] if cap["end"] > 0 else cap["start"] + 5.0,
                    label=cap["text"],
                    tier="caption",
                    metadata={"index": i},
                )
                for i, cap in enumerate(captions)
            ]
        )

        # Combined translation (all captions)
        translation = " ".join(cap["text"] for cap in captions)

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=[],  # Caption-based, no glosses
            text=translation,
            dataset=self.NAME,
            split=self.split,
            metadata={
                "language": self.language,
                "num_captions": len(captions),
                **meta,
            },
        )

    def _find_video(self, sample_id: str) -> Path | None:
        """Find video file."""
        search_dirs = [
            self.root / "videos",
            self.root / self.language / "videos",
            self.root,
        ]

        for search_dir in search_dirs:
            if not search_dir.exists():
                continue

            for ext in [".mp4", ".webm", ".avi", ".mkv"]:
                video_path = search_dir / f"{sample_id}{ext}"
                if video_path.exists():
                    return video_path

        return None

    def get_caption_vocabulary(self, min_freq: int = 1) -> dict[str, int]:
        """Get word frequencies from captions."""
        word_freq: dict[str, int] = {}

        for captions in self._captions.values():
            for cap in captions:
                words = cap["text"].lower().split()
                for word in words:
                    # Clean punctuation
                    word = re.sub(r"[^\w\s]", "", word)
                    if word:
                        word_freq[word] = word_freq.get(word, 0) + 1

        return {w: f for w, f in word_freq.items() if f >= min_freq}

    def summary(self) -> dict[str, Any]:
        """Get dataset summary."""
        base = super().summary()
        base.update(
            {
                "language": self.language,
                "num_videos": len(self._captions),
                "total_captions": sum(len(c) for c in self._captions.values()),
                "annotation_type": "caption",
            }
        )
        return base


class YTASLDataset(BaseDataset):
    """
    YouTube ASL (YT-ASL) dataset loader.

    Large-scale ASL dataset from YouTube with subtitle annotations.
    """

    NAME = "yt_asl"
    TASKS = ["sign_language_translation", "caption_based_recognition"]
    LANGUAGES = ["asl"]
    CITATION = """
    @article{uthus2023youtube,
        title={YouTube-ASL: A Large-Scale, Open-Domain American Sign Language-English Parallel Corpus},
        author={Uthus, David and others},
        year={2023}
    }
    """

    def __init__(
        self,
        root: str | Path,
        split: str = "train",
        feature_root: str | Path | None = None,
        load_poses: bool = False,
        pose_format: str = "mediapipe",
        fps: float = 25.0,
    ):
        """
        Initialize YT-ASL dataset.

        Args:
            root: Path to YT-ASL directory
            split: Data split ('train', 'val', 'test')
            feature_root: Path to pre-extracted features
            load_poses: Whether to load pose data
            pose_format: Pose format to load
            fps: Video FPS
        """
        self._root = Path(root)
        self._subtitles = self._load_subtitles()
        self._manifest = self._load_manifest()

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

    def _load_subtitles(self) -> dict[str, list[dict[str, Any]]]:
        """Load subtitle files (VTT/SRT)."""
        subtitles = {}

        sub_dirs = [
            self._root / "subtitles",
            self._root / "captions",
            self._root,
        ]

        for sub_dir in sub_dirs:
            if not sub_dir.exists():
                continue

            # Load VTT files
            for vtt_file in sub_dir.glob("*.vtt"):
                video_id = vtt_file.stem
                subtitles[video_id] = self._parse_vtt(vtt_file)

            # Load SRT files
            for srt_file in sub_dir.glob("*.srt"):
                video_id = srt_file.stem
                if video_id not in subtitles:
                    subtitles[video_id] = self._parse_srt(srt_file)

            if subtitles:
                break

        return subtitles

    def _parse_vtt(self, vtt_path: Path) -> list[dict[str, Any]]:
        """Parse WebVTT subtitle file."""
        cues: list[dict[str, Any]] = []

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

        # Skip header
        lines = content.strip().split("\n")
        i = 0
        while i < len(lines) and not _TIME_CHECK_PATTERN.match(lines[i]):
            i += 1

        current_cue: dict[str, Any] | None = None
        while i < len(lines):
            line = lines[i].strip()

            # Time line
            time_match = _VTT_TIME_PATTERN.match(line)
            if time_match:
                if current_cue:
                    cues.append(current_cue)
                start = self._parse_vtt_time(time_match.group(1))
                end = self._parse_vtt_time(time_match.group(2))
                current_cue = {"start": start, "end": end, "text": ""}
            elif current_cue and line and not line.isdigit():
                # Text line
                if current_cue["text"]:
                    current_cue["text"] += " "
                current_cue["text"] += line

            i += 1

        if current_cue:
            cues.append(current_cue)

        return cues

    def _parse_vtt_time(self, time_str: str) -> float:
        """Parse VTT time format."""
        parts = time_str.replace(",", ".").split(":")
        if len(parts) == 3:
            h, m, s = parts
            return int(h) * 3600 + int(m) * 60 + float(s)
        elif len(parts) == 2:
            m, s = parts
            return int(m) * 60 + float(s)
        return 0.0

    def _parse_srt(self, srt_path: Path) -> list[dict[str, Any]]:
        """Parse SRT subtitle file."""
        cues = []

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

        blocks = re.split(r"\n\n+", content.strip())

        for block in blocks:
            lines = block.strip().split("\n")
            if len(lines) < 2:
                continue

            # Find time line
            time_match = None
            text_start = 0
            for j, line in enumerate(lines):
                time_match = _SRT_TIME_PATTERN.match(line)
                if time_match:
                    text_start = j + 1
                    break

            if not time_match:
                continue

            start = self._parse_vtt_time(time_match.group(1))
            end = self._parse_vtt_time(time_match.group(2))
            text = " ".join(lines[text_start:])

            # Clean HTML tags
            text = re.sub(r"<[^>]+>", "", text)

            cues.append({"start": start, "end": end, "text": text})

        return cues

    def _load_manifest(self) -> dict[str, dict[str, Any]]:
        """Load manifest JSON."""
        manifest = {}

        manifest_paths = [
            self._root / "manifest.json",
            self._root / "metadata.json",
        ]

        for path in manifest_paths:
            if path.exists():
                with open(path, encoding="utf-8") as f:
                    data = json.load(f)
                    if isinstance(data, list):
                        for item in data:
                            vid_id = item.get("video_id", "") or item.get("id", "")
                            manifest[vid_id] = item
                    else:
                        manifest = data
                break

        return manifest

    def _load_index(self) -> list[str]:
        """Get sample IDs."""
        video_ids = []

        for vid_id in self._subtitles.keys():
            meta = self._manifest.get(vid_id, {})
            vid_split = meta.get("split", "train")

            if vid_split == self.split or self.split == "all":
                video_ids.append(vid_id)

        return video_ids

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single YT-ASL sample."""
        subtitles = self._subtitles.get(sample_id, [])
        meta = self._manifest.get(sample_id, {})

        # Find video
        video_path = self._find_video(sample_id)

        # Create segments from subtitles
        segments = SegmentList(
            [
                Segment(
                    start=sub["start"],
                    end=sub["end"],
                    label=sub["text"],
                    tier="subtitle",
                    metadata={"index": i},
                )
                for i, sub in enumerate(subtitles)
            ]
        )

        # Combined text
        translation = " ".join(sub["text"] for sub in subtitles)

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=[],  # No glosses for subtitle-based dataset
            text=translation,
            dataset=self.NAME,
            split=self.split,
            metadata={
                "num_subtitles": len(subtitles),
                **meta,
            },
        )

    def _find_video(self, sample_id: str) -> Path | None:
        """Find video file."""
        video_dirs = [
            self.root / "videos",
            self.root,
        ]

        for video_dir in video_dirs:
            if not video_dir.exists():
                continue

            for ext in [".mp4", ".webm", ".avi", ".mkv"]:
                video_path = video_dir / f"{sample_id}{ext}"
                if video_path.exists():
                    return video_path

        return None

    def summary(self) -> dict[str, Any]:
        """Get dataset summary."""
        base = super().summary()
        base.update(
            {
                "num_videos": len(self._subtitles),
                "total_subtitles": sum(len(s) for s in self._subtitles.values()),
                "annotation_type": "subtitle",
            }
        )
        return base
