"""
WLASL (Word-Level American Sign Language) dataset loader.

WLASL is a large-scale word-level American Sign Language dataset
containing over 2000 signs with multiple videos per sign.

Dataset structure:
    /path/to/wlasl/
        WLASL_v0.3.json           # Main annotation file
        videos/                    # Video files
            *.mp4
        nslt_*.json               # Various split files

Reference:
    Li, D., et al. "Word-level Deep Sign Language Recognition from Video:
    A New Large-scale Dataset and Methods Comparison." WACV 2020.

Usage:
    from sltk.data.datasets import WLASLDataset

    dataset = WLASLDataset(
        root="/path/to/wlasl",
        split="train",
        num_classes=100,  # WLASL100, WLASL300, WLASL1000, WLASL2000
    )

    sample = dataset[0]
    print(sample.glosses)  # e.g., ['book']
"""

import json
from pathlib import Path
from typing import Any

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


class WLASLDataset(BaseDataset):
    """
    WLASL dataset loader for word-level American Sign Language recognition.

    Supports different vocabulary sizes: WLASL100, WLASL300, WLASL1000, WLASL2000.
    """

    NAME = "wlasl"
    TASKS = ["isolated_sign_recognition"]
    LANGUAGES = ["asl"]
    CITATION = """
    @inproceedings{li2020word,
        title={Word-level Deep Sign Language Recognition from Video},
        author={Li, Dongxu and Rodriguez, Cristian and Yu, Xin and Li, Hongdong},
        booktitle={WACV},
        year={2020}
    }
    """

    # Standard WLASL subsets
    SUBSETS = {100, 300, 1000, 2000}

    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,
        num_classes: int = 2000,
        annotation_file: str = "WLASL_v0.3.json",
    ):
        """
        Initialize WLASL dataset.

        Args:
            root: Path to WLASL 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 ('mediapipe', 'wilor', etc.)
            fps: Video FPS
            num_classes: Vocabulary size (100, 300, 1000, 2000)
            annotation_file: Name of annotation JSON file
        """
        self.num_classes = num_classes
        self.annotation_file = annotation_file

        # Load annotations before parent init
        self._root = Path(root)
        self._annotations = self._load_annotations()
        self._gloss_to_idx = self._build_gloss_index()
        self._samples = self._build_samples()

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

    def _load_annotations(self) -> list[dict[str, Any]]:
        """Load WLASL annotation JSON."""
        anno_path = self._root / self.annotation_file
        if not anno_path.exists():
            # Try alternative locations
            for alt in ["wlasl_class_list.txt", "nslt_2000.json"]:
                alt_path = self._root / alt
                if alt_path.exists():
                    return self._load_alternative_format(alt_path)
            return []

        with open(anno_path, encoding="utf-8") as f:
            return json.load(f)

    def _load_alternative_format(self, path: Path) -> list[dict[str, Any]]:
        """Load alternative annotation formats."""
        if path.suffix == ".txt":
            # Class list format
            glosses = []
            with open(path, encoding="utf-8") as f:
                for line in f:
                    gloss = line.strip()
                    if gloss:
                        glosses.append({"gloss": gloss, "instances": []})
            return glosses
        elif path.suffix == ".json":
            with open(path, encoding="utf-8") as f:
                return json.load(f)
        return []

    def _build_gloss_index(self) -> dict[str, int]:
        """Build mapping from gloss to class index."""
        index = {}
        for i, entry in enumerate(self._annotations):
            if i >= self.num_classes:
                break
            gloss = entry.get("gloss", "")
            index[gloss.lower()] = i
        return index

    def _build_samples(self) -> list[dict[str, Any]]:
        """Build list of all samples with their metadata."""
        samples = []

        for gloss_idx, entry in enumerate(self._annotations):
            if gloss_idx >= self.num_classes:
                break

            gloss = entry.get("gloss", "")

            for instance in entry.get("instances", []):
                video_id = instance.get("video_id", "")
                split = instance.get("split", "train")

                samples.append(
                    {
                        "id": f"{gloss}_{video_id}",
                        "video_id": video_id,
                        "gloss": gloss,
                        "gloss_idx": gloss_idx,
                        "split": split,
                        "frame_start": instance.get("frame_start", 0),
                        "frame_end": instance.get("frame_end", -1),
                        "fps": instance.get("fps", 25),
                        "signer_id": instance.get("signer_id", ""),
                        "bbox": instance.get("bbox", None),
                        "source": instance.get("source", ""),
                    }
                )

        return samples

    def _load_index(self) -> list[str]:
        """Get sample IDs for current split."""
        return [s["id"] for s in self._samples if s["split"] == self.split or self.split == "all"]

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single WLASL sample."""
        # Find sample by ID
        sample_meta = None
        for s in self._samples:
            if s["id"] == sample_id:
                sample_meta = s
                break

        if sample_meta is None:
            raise SampleNotFoundError(sample_id, self.NAME)

        gloss = sample_meta["gloss"]

        # Find video file
        video_path = self._find_video(sample_meta["video_id"])

        # Create segment for isolated sign
        segments = SegmentList(
            [
                Segment(
                    start=sample_meta["frame_start"] / sample_meta["fps"],
                    end=(
                        sample_meta["frame_end"] / sample_meta["fps"]
                        if sample_meta["frame_end"] > 0
                        else float("inf")
                    ),
                    label=gloss,
                    tier="gloss",
                    metadata={
                        "gloss_idx": sample_meta["gloss_idx"],
                        "frame_start": sample_meta["frame_start"],
                        "frame_end": sample_meta["frame_end"],
                    },
                )
            ]
        )

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=[gloss],
            dataset=self.NAME,
            split=self.split,
            metadata={
                "video_id": sample_meta["video_id"],
                "gloss_idx": sample_meta["gloss_idx"],
                "signer_id": sample_meta["signer_id"],
                "bbox": sample_meta["bbox"],
                "source": sample_meta["source"],
                "num_classes": self.num_classes,
            },
        )

    def _find_video(self, video_id: str) -> Path | None:
        """Find video file by ID."""
        video_dir = self.root / "videos"
        if not video_dir.exists():
            video_dir = self.root

        # Try different extensions and patterns
        for ext in [".mp4", ".avi", ".mov", ".webm"]:
            patterns = [
                video_dir / f"{video_id}{ext}",
                video_dir / f"{video_id.zfill(5)}{ext}",
            ]
            for pattern in patterns:
                if pattern.exists():
                    return pattern

        return None

    def get_vocabulary(self) -> list[str]:
        """Get all glosses in the vocabulary."""
        return [e["gloss"] for e in self._annotations[: self.num_classes]]

    def get_class_samples(self, gloss: str) -> list[str]:
        """Get all sample IDs for a specific gloss."""
        gloss_lower = gloss.lower()
        return [
            s["id"]
            for s in self._samples
            if s["gloss"].lower() == gloss_lower and s["split"] == self.split
        ]

    def get_signer_samples(self, signer_id: str) -> list[str]:
        """Get all sample IDs from a specific signer."""
        return [
            s["id"]
            for s in self._samples
            if s["signer_id"] == signer_id and s["split"] == self.split
        ]

    def summary(self) -> dict[str, Any]:
        """Get dataset summary."""
        base = super().summary()
        base.update(
            {
                "num_classes": self.num_classes,
                "vocabulary_size": len(self.get_vocabulary()),
                "total_samples": len(self._samples),
                "split_samples": len(self._index),
            }
        )
        return base
