"""
ASL Citizen dataset loader.

ASL Citizen is a community-sourced ASL dataset collected through
a mobile app, with videos from diverse signers.

Dataset structure:
    /path/to/asl-citizen/
        videos/
            *.mp4
        metadata.json or asl_citizen.csv
        splits/
            train.txt
            val.txt
            test.txt

Reference:
    Desai, A., et al. "ASL Citizen: A Community-Sourced Dataset
    for Advancing Isolated Sign Language Recognition." 2023.

Usage:
    from sltk.data.datasets import ASLCitizenDataset

    dataset = ASLCitizenDataset(
        root="/path/to/asl-citizen",
        split="train",
    )

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

import csv
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 ASLCitizenDataset(BaseDataset):
    """
    ASL Citizen dataset loader.

    Community-sourced isolated ASL signs from diverse signers.
    """

    NAME = "asl_citizen"
    TASKS = ["isolated_sign_recognition"]
    LANGUAGES = ["asl"]
    CITATION = """
    @article{desai2023asl,
        title={ASL Citizen: A Community-Sourced Dataset for Advancing Isolated Sign Language Recognition},
        author={Desai, Aashaka and Berger, Lauren and Minakov, Fyodor and Huenerfauth, Matt and Veeramachaneni, Kalyan and Yin, Naomi 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 = 30.0,
    ):
        """
        Initialize ASL Citizen dataset.

        Args:
            root: Path to ASL Citizen 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._metadata = self._load_metadata()
        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_metadata(self) -> dict[str, dict[str, Any]]:
        """Load metadata file."""
        metadata = {}

        # Try JSON format
        json_path = self._root / "metadata.json"
        if json_path.exists():
            with open(json_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
            return metadata

        # Try CSV format
        csv_paths = [
            self._root / "asl_citizen.csv",
            self._root / "metadata.csv",
            self._root / "annotations.csv",
        ]

        for csv_path in csv_paths:
            if csv_path.exists():
                with open(csv_path, encoding="utf-8") as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        vid_id = (
                            row.get("video_id", "") or row.get("id", "") or row.get("filename", "")
                        )
                        if vid_id:
                            # Remove extension if present
                            vid_id = Path(vid_id).stem
                            metadata[vid_id] = {
                                "gloss": row.get("label", "") or row.get("gloss", ""),
                                "signer_id": row.get("participant_id", "")
                                or row.get("signer_id", ""),
                                "split": row.get("split", ""),
                            }
                break

        return metadata

    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", "val", "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] = [Path(line.strip()).stem for line in f if line.strip()]

        return splits

    def _load_index(self) -> list[str]:
        """Load sample IDs for current split."""
        # If split files exist, use them
        if self.split in self._split_files:
            return self._split_files[self.split]

        # Otherwise, filter metadata by split field
        if self._metadata:
            return [
                vid_id
                for vid_id, meta in self._metadata.items()
                if meta.get("split", "") == self.split or self.split == "all"
            ]

        # Fallback: scan video directory
        video_dir = self.root / "videos"
        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 ASL Citizen sample."""
        meta = self._metadata.get(sample_id, {})

        gloss = meta.get("gloss", "") or meta.get("label", "")
        signer_id = meta.get("signer_id", "") or meta.get("participant_id", "")

        # Find video
        video_path = self._find_video(sample_id)

        # Create segment
        segments = SegmentList(
            [
                Segment(
                    start=0.0,
                    end=float("inf"),
                    label=gloss,
                    tier="gloss",
                )
            ]
        )

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=[gloss] if gloss else [],
            dataset=self.NAME,
            split=self.split,
            metadata={
                "signer_id": signer_id,
                **{
                    k: v
                    for k, v in meta.items()
                    if k not in ["gloss", "label", "signer_id", "participant_id"]
                },
            },
        )

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

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

        return None

    def get_vocabulary(self) -> list[str]:
        """Get all unique glosses."""
        vocab = set()
        for meta in self._metadata.values():
            gloss = meta.get("gloss", "") or meta.get("label", "")
            if gloss:
                vocab.add(gloss)
        return sorted(vocab)

    def get_signers(self) -> list[str]:
        """Get all unique signer IDs."""
        signers = set()
        for meta in self._metadata.values():
            signer = meta.get("signer_id", "") or meta.get("participant_id", "")
            if signer:
                signers.add(signer)
        return sorted(signers)

    def get_samples_by_signer(self, signer_id: str) -> list[str]:
        """Get sample IDs for a specific signer."""
        return [
            vid_id
            for vid_id, meta in self._metadata.items()
            if (
                meta.get("signer_id", "") == signer_id
                or meta.get("participant_id", "") == signer_id
            )
            and vid_id in self._index
        ]

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