"""
Phoenix-2014T dataset loader.

Phoenix-2014T is a German Sign Language dataset from weather
forecast broadcasts with both gloss and German text annotations.

Dataset structure:
    /path/to/phoenix-2014t/
        annotations/manual/
            train.corpus.csv
            dev.corpus.csv
            test.corpus.csv
        features/fullFrame-210x260px/
            train/
            dev/
            test/

Reference:
    Camgoz, N.C., et al. "Neural Sign Language Translation."
    CVPR 2018.

Usage:
    from sltk.data.datasets import Phoenix2014TDataset

    dataset = Phoenix2014TDataset(
        root="/path/to/phoenix-2014t",
        split="train",
    )

    sample = dataset[0]
    print(sample.glosses)  # ['WOLKEN', 'SONNE', ...]
    print(sample.translation)  # German 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 Phoenix2014TDataset(BaseDataset):
    """
    Phoenix-2014T dataset loader for German Sign Language.

    Provides both gloss-level and translation-level annotations.
    """

    NAME = "phoenix2014t"
    TASKS = ["sign_language_translation", "continuous_sign_recognition"]
    LANGUAGES = ["dgs"]  # German Sign Language
    CITATION = """
    @inproceedings{camgoz2018neural,
        title={Neural Sign Language Translation},
        author={Camgoz, Necati Cihan and Hadfield, Simon and Koller, Oscar and Ney, Hermann and Bowden, Richard},
        booktitle={CVPR},
        year={2018}
    }
    """

    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,
        use_text: bool = True,
    ):
        """
        Initialize Phoenix-2014T dataset.

        Args:
            root: Path to Phoenix-2014T 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
            use_text: Whether to include German text translations
        """
        self.use_text = use_text

        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 files for all splits."""
        annotations = {}

        for split_name in ["train", "dev", "test"]:
            csv_paths = [
                self._root / "annotations" / "manual" / f"{split_name}.corpus.csv",
                self._root / f"{split_name}.corpus.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 Phoenix annotation CSV."""
        samples = []

        with open(csv_path, encoding="utf-8") as f:
            # Phoenix uses | as delimiter
            reader = csv.DictReader(f, delimiter="|")

            for row in reader:
                sample_id = row.get("name", "") or row.get("id", "")
                gloss_str = row.get("annotation", "") or row.get("orth", "")
                translation = row.get("translation", "") or row.get("text", "")

                # Parse glosses (space-separated)
                glosses = gloss_str.strip().split() if gloss_str else []

                samples.append(
                    {
                        "id": sample_id,
                        "glosses": glosses,
                        "translation": translation,
                        "signer_id": row.get("speaker", ""),
                        "video_folder": row.get("video", ""),
                    }
                )

        return samples

    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]:
        """Get sample IDs for current split."""
        split_name = self._get_split_name()
        samples = self._annotations.get(split_name, [])
        return [s["id"] for s in samples if s["id"]]

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single Phoenix-2014T sample."""
        split_name = self._get_split_name()
        samples = self._annotations.get(split_name, [])

        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/frames
        video_path = self._find_video(sample_id, sample_meta.get("video_folder", ""), split_name)

        glosses = sample_meta["glosses"]

        # Create segments from glosses
        # Phoenix doesn't have timing, so create placeholder segments
        segments = SegmentList(
            [
                Segment(
                    start=0.0,
                    end=float("inf"),
                    label=" ".join(glosses),
                    tier="gloss",
                )
            ]
        )

        translation = sample_meta["translation"] if self.use_text else ""

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=None,
            segments=segments,
            glosses=glosses,
            text=translation,
            signer_id=sample_meta["signer_id"],
            dataset=self.NAME,
            split=self.split,
            metadata={
                "video_folder": sample_meta["video_folder"],
            },
        )

    def _find_video(self, sample_id: str, video_folder: str, split_name: str) -> Path | None:
        """Find video or frame directory."""
        # Phoenix often uses frame directories instead of videos
        search_paths = [
            self.root / "features" / "fullFrame-210x260px" / split_name / video_folder,
            self.root / "features" / "fullFrame-210x260px" / split_name / sample_id,
            self.root / split_name / video_folder,
            self.root / split_name / sample_id,
        ]

        # Try frame directories first
        for path in search_paths:
            if path.exists() and path.is_dir():
                return path

        # Try video files
        for base_dir in [self.root / split_name, self.root / "videos" / split_name]:
            if not base_dir.exists():
                continue
            for ext in [".mp4", ".avi"]:
                video_path = base_dir / f"{sample_id}{ext}"
                if video_path.exists():
                    return video_path

        return None

    def get_vocabulary(self) -> list[str]:
        """Get all unique glosses in the dataset."""
        vocab = set()
        for samples in self._annotations.values():
            for s in samples:
                vocab.update(s["glosses"])
        return sorted(vocab)

    def get_signers(self) -> list[str]:
        """Get all 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 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()),
                "has_translation": self.use_text,
            }
        )
        return base
