"""
CSL-Daily (Chinese Sign Language Daily) dataset loader.

CSL-Daily is a Chinese Sign Language dataset containing daily
conversation scenarios with sentence-level annotations.

Dataset structure:
    /path/to/csl-daily/
        sentence/
            train/
                *.mp4
            dev/
                *.mp4
            test/
                *.mp4
        sentence_label.json
        splits/
            train.txt
            dev.txt
            test.txt

Reference:
    Zhou, H., et al. "Improving Sign Language Translation with
    Monolingual Data by Sign Back-Translation." CVPR 2021.

Usage:
    from sltk.data.datasets import CSLDailyDataset

    dataset = CSLDailyDataset(
        root="/path/to/csl-daily",
        split="train",
    )

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

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


class CSLDailyDataset(BaseDataset):
    """
    CSL-Daily dataset loader for Chinese Sign Language.

    Provides sentence-level sign language translation pairs.
    """

    NAME = "csl_daily"
    TASKS = ["sign_language_translation", "continuous_sign_recognition"]
    LANGUAGES = ["csl"]
    CITATION = """
    @inproceedings{zhou2021improving,
        title={Improving Sign Language Translation with Monolingual Data by Sign Back-Translation},
        author={Zhou, Hao and Zhou, Wengang and Qi, Weizhen and Pu, Junfu and Li, Houqiang},
        booktitle={CVPR},
        year={2021}
    }
    """

    # Split mappings
    SPLIT_MAPPING = {
        "train": "train",
        "val": "dev",
        "dev": "dev",
        "test": "test",
    }

    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 CSL-Daily dataset.

        Args:
            root: Path to CSL-Daily directory
            split: Data split ('train', 'val'/'dev', '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._labels = self._load_labels()
        self._split_files = self._load_split_files()

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

    def _load_labels(self) -> dict[str, dict[str, Any]]:
        """Load sentence labels."""
        labels = {}

        # Try different label file locations
        label_paths = [
            self._root / "sentence_label.json",
            self._root / "labels.json",
            self._root / "annotations.json",
        ]

        for path in label_paths:
            if path.exists():
                with open(path, encoding="utf-8") as f:
                    data = json.load(f)
                    # Handle different formats
                    if isinstance(data, dict):
                        labels = data
                    elif isinstance(data, list):
                        for item in data:
                            if "name" in item:
                                labels[item["name"]] = item
                            elif "id" in item:
                                labels[item["id"]] = item
                break

        return labels

    def _load_split_files(self) -> dict[str, list[str]]:
        """Load split file lists."""
        splits = {}

        split_dir = self._root / "splits"
        if not split_dir.exists():
            split_dir = self._root

        for split_name in ["train", "dev", "test"]:
            split_file = split_dir / f"{split_name}.txt"
            if split_file.exists():
                with open(split_file, encoding="utf-8") as f:
                    splits[split_name] = [line.strip() for line in f if line.strip()]

        return splits

    def _get_split_name(self) -> str:
        """Get internal split name."""
        return self.SPLIT_MAPPING.get(self.split, self.split)

    def _load_index(self) -> list[str]:
        """Load sample IDs for current split."""
        split_name = self._get_split_name()

        # If split files exist, use them
        if split_name in self._split_files:
            return self._split_files[split_name]

        # Otherwise, scan video directory
        video_dir = self._root / "sentence" / split_name
        if not video_dir.exists():
            video_dir = self._root / split_name

        if video_dir.exists():
            return [p.stem for p in video_dir.glob("*.mp4")]

        return []

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single CSL-Daily sample."""
        split_name = self._get_split_name()

        # Find video
        video_path = self._find_video(sample_id, split_name)

        # Get label info
        label_info = self._labels.get(sample_id, {})

        # Chinese sentence (translation)
        sentence = label_info.get("text", "")
        if not sentence:
            sentence = label_info.get("sentence", "")
            if not sentence:
                sentence = label_info.get("label", "")

        # Glosses (if available)
        glosses = label_info.get("gloss", [])
        if isinstance(glosses, str):
            glosses = glosses.split()

        # Create segments from glosses if available
        if glosses:
            segments = SegmentList(
                [
                    Segment(
                        start=0.0,
                        end=float("inf"),
                        label=" ".join(glosses) if isinstance(glosses, list) else glosses,
                        tier="gloss",
                    )
                ]
            )
        else:
            segments = SegmentList()

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=glosses if isinstance(glosses, list) else [],
            text=sentence,
            signer_id=label_info.get("signer_id", label_info.get("signer", "")),
            dataset=self.NAME,
            split=self.split,
            metadata={
                "sentence": sentence,
            },
        )

    def _find_video(self, sample_id: str, split_name: str) -> Path | None:
        """Find video file."""
        # Try different locations
        candidates = [
            self.root / "sentence" / split_name / f"{sample_id}.mp4",
            self.root / split_name / f"{sample_id}.mp4",
            self.root / "videos" / f"{sample_id}.mp4",
            self.root / f"{sample_id}.mp4",
        ]

        for path in candidates:
            if path.exists():
                return path

        return None

    def get_vocabulary(self) -> list[str]:
        """Get all unique glosses."""
        vocab = set()
        for label_info in self._labels.values():
            glosses = label_info.get("gloss", [])
            if isinstance(glosses, str):
                glosses = glosses.split()
            vocab.update(glosses)
        return sorted(vocab)

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