"""
How2Sign dataset loader.

How2Sign is a large-scale multimodal and multilingual dataset for
American Sign Language with aligned transcriptions.

Dataset structure:
    /path/to/how2sign/
        video_level/
            train/
                *.mp4
            val/
                *.mp4
            test/
                *.mp4
        sentence_level/
            train/
                *.mp4
            val/
                *.mp4
            test/
                *.mp4
        how2sign_realigned_train.csv
        how2sign_realigned_val.csv
        how2sign_realigned_test.csv

Reference:
    Duarte, A., et al. "How2Sign: A Large-scale Multimodal Dataset
    for Continuous American Sign Language." CVPR 2021.

Usage:
    from sltk.data.datasets import How2SignDataset

    dataset = How2SignDataset(
        root="/path/to/how2sign",
        split="train",
        level="sentence",  # 'video' or 'sentence'
    )

    sample = dataset[0]
    print(sample.translation)  # English sentence
"""

import csv
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 How2SignDataset(BaseDataset):
    """
    How2Sign dataset loader for continuous ASL.

    Supports both video-level and sentence-level samples.
    """

    NAME = "how2sign"
    TASKS = ["sign_language_translation", "continuous_sign_recognition"]
    LANGUAGES = ["asl"]
    CITATION = """
    @inproceedings{duarte2021how2sign,
        title={How2Sign: A Large-scale Multimodal Dataset for Continuous American Sign Language},
        author={Duarte, Amanda and Palaskar, Shruti and Ventura, Lucas and Ghadiyaram, Deepti and DeHaan, Kenneth and Metze, Florian and Torres, Jordi and Giro-i-Nieto, Xavier},
        booktitle={CVPR},
        year={2021}
    }
    """

    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 = 24.0,
        level: str = "sentence",
        view: str = "front",
    ):
        """
        Initialize How2Sign dataset.

        Args:
            root: Path to How2Sign 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 (default 24 for How2Sign)
            level: Granularity level ('video' or 'sentence')
            view: Camera view ('front', 'side', etc.)
        """
        self.level = level
        self.view = view

        self._root = Path(root)
        self._annotations = self._load_annotations()

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

    def _load_annotations(self) -> dict[str, list[dict[str, Any]]]:
        """Load annotation CSV files for all splits."""
        annotations = {}

        for split_name in ["train", "val", "test"]:
            csv_paths = [
                self._root / f"how2sign_realigned_{split_name}.csv",
                self._root / f"how2sign_{split_name}.csv",
                self._root / f"{split_name}.csv",
            ]

            for csv_path in csv_paths:
                if csv_path.exists():
                    annotations[split_name] = self._parse_csv(csv_path)
                    break

            if split_name not in annotations:
                annotations[split_name] = []

        return annotations

    def _parse_csv(self, csv_path: Path) -> list[dict[str, Any]]:
        """Parse annotation CSV file."""
        samples = []

        with open(csv_path, encoding="utf-8", errors="replace") as f:
            # Detect delimiter
            sample = f.read(1024)
            f.seek(0)

            if "\t" in sample:
                delimiter = "\t"
            else:
                delimiter = ","

            reader = csv.DictReader(f, delimiter=delimiter)

            for row in reader:
                # Handle different column naming conventions
                sample_id = (
                    row.get("SENTENCE_NAME", "")
                    or row.get("sentence_name", "")
                    or row.get("VIDEO_NAME", "")
                    or row.get("video_name", "")
                    or row.get("id", "")
                )

                sentence = (
                    row.get("SENTENCE", "")
                    or row.get("sentence", "")
                    or row.get("TEXT", "")
                    or row.get("text", "")
                )

                start_time = float(row.get("START_REALIGNED", 0) or row.get("start", 0) or 0)
                end_time = float(row.get("END_REALIGNED", 0) or row.get("end", 0) or 0)

                samples.append(
                    {
                        "id": sample_id,
                        "sentence": sentence,
                        "start": start_time,
                        "end": end_time,
                        "video_id": row.get("VIDEO_ID", "") or row.get("video_id", ""),
                        "signer_id": row.get("SIGNER", "") or row.get("signer", ""),
                    }
                )

        return samples

    def _load_index(self) -> list[str]:
        """Get sample IDs for current split."""
        samples = self._annotations.get(self.split, [])
        return [s["id"] for s in samples if s["id"]]

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single How2Sign sample."""
        # Find sample in annotations
        samples = self._annotations.get(self.split, [])
        sample_meta = None
        for s in samples:
            if s["id"] == sample_id:
                sample_meta = s
                break

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

        # Find video file
        video_path = self._find_video(sample_id)

        # Create segment
        segments = SegmentList(
            [
                Segment(
                    start=sample_meta["start"],
                    end=sample_meta["end"] if sample_meta["end"] > 0 else float("inf"),
                    label=sample_meta["sentence"],
                    tier="translation",
                )
            ]
        )

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=[],  # How2Sign doesn't have gloss annotations
            text=sample_meta["sentence"],
            signer_id=sample_meta["signer_id"],
            dataset=self.NAME,
            split=self.split,
            metadata={
                "video_id": sample_meta["video_id"],
                "start_time": sample_meta["start"],
                "end_time": sample_meta["end"],
                "level": self.level,
                "view": self.view,
            },
        )

    def _find_video(self, sample_id: str) -> Path | None:
        """Find video file."""
        # Try different directory structures
        search_dirs = [
            self.root / f"{self.level}_level" / self.split,
            self.root / self.level / self.split,
            self.root / self.split,
            self.root / "videos" / self.split,
            self.root / "videos",
        ]

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

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

        return None

    def get_signers(self) -> list[str]:
        """Get list of unique signer IDs."""
        signers = set()
        for samples in self._annotations.values():
            for s in samples:
                if s.get("signer_id"):
                    signers.add(s["signer_id"])
        return sorted(signers)

    def get_samples_by_signer(self, signer_id: str) -> list[str]:
        """Get sample IDs for a specific signer."""
        samples = self._annotations.get(self.split, [])
        return [s["id"] for s in samples if s.get("signer_id") == signer_id]

    def summary(self) -> dict[str, Any]:
        """Get dataset summary."""
        base = super().summary()
        base.update(
            {
                "level": self.level,
                "view": self.view,
                "num_signers": len(self.get_signers()),
                "total_samples": sum(len(s) for s in self._annotations.values()),
            }
        )
        return base
