"""
BSLCP (BSL Corpus Project) dataset loader.

BSLCP is a multiview British Sign Language corpus with ELAN annotations.
Videos are recorded from multiple camera angles with detailed gloss annotations.

Dataset structure:
    {SLTK_DATASETS_ROOT}/multiview/BSLCP/
        BF01n.mov, BF01n.eaf   # Signer BF, session 01, narration
        BF01c.mov, BF01c.eaf   # Signer BF, session 01, conversation
        ...

Features available at:
    {SLTK_FEATURE_POOL_ROOT}/features2/BSLCP2/
        wilor/BF01n_wilor.h5
        smpl/BF01n_nlf.h5

Usage:
    from sltk.data.datasets import BSLCPDataset

    # Uses default paths from environment or hardcoded defaults
    dataset = BSLCPDataset(load_poses=True, pose_format="wilor")

    # Or specify paths explicitly
    dataset = BSLCPDataset(
        root="/path/to/BSLCP",
        feature_root="/path/to/features/BSLCP2",
        load_poses=True,
        pose_format="wilor",
    )

    sample = dataset[0]
    print(sample.segments)  # Gloss annotations
    print(sample.poses)     # WiLoR hand poses
"""

import re
from pathlib import Path

from sltk.data.core import PoseSequence, Sample, SegmentList
from sltk.data.datasets.base import BaseDataset
from sltk.exceptions import SampleNotFoundError
from sltk.io.elan import get_eaf_tiers, read_eaf


def _get_default_video_root() -> Path:
    """Get default video root from config."""
    import os

    # Use environment variable or hardcoded path to scratch storage
    custom_root = os.environ.get("SLTK_BSLCP_ROOT")
    if custom_root:
        return Path(custom_root)

    # Hardcoded path to BSLCP data
    scratch_path = Path("/mnt/fast/nobackup/scratch4weeks/ef0036/videos/bslcp/BSLCP")
    if scratch_path.exists():
        return scratch_path

    # Fallback to research network path
    from sltk.config import get_default_dataset_roots

    return get_default_dataset_roots()["multiview"] / "BSLCP"


def _get_default_feature_root() -> Path:
    """Get default feature root from config."""
    from sltk.config import _get_feature_pool_root

    return _get_feature_pool_root() / "features2" / "bslcp"


class BSLCPDataset(BaseDataset):
    """
    BSLCP dataset loader.

    Loads gloss annotations from ELAN files and optionally links
    to pre-extracted pose features (WiLoR, SMPL/NLF).
    """

    NAME = "bslcp"
    TASKS = ["gloss_recognition", "sign_spotting"]
    LANGUAGES = ["bsl"]
    CITATION = "Schembri et al. (2013) Building the British Sign Language Corpus"

    # These are now computed dynamically to respect environment variables
    @property
    def DEFAULT_VIDEO_ROOT(self) -> Path:  # type: ignore[override]
        return _get_default_video_root()

    @property
    def DEFAULT_FEATURE_ROOT(self) -> Path:  # type: ignore[override]
        return _get_default_feature_root()

    # Common gloss tier names in BSLCP
    GLOSS_TIERS = [
        "RH-IDgloss",  # Right hand ID gloss
        "LH-IDgloss",  # Left hand ID gloss
        "RH-Gloss",  # Right hand gloss
        "LH-Gloss",  # Left hand gloss
    ]

    def __init__(
        self,
        root: str | Path | None = None,
        split: str = "all",
        feature_root: str | Path | None = None,
        load_poses: bool = False,
        pose_format: str = "wilor",
        fps: float = 25.0,
        gloss_tiers: list[str] | None = None,
        signer_filter: list[str] | None = None,
    ):
        """
        Initialize BSLCP dataset.

        Args:
            root: Path to BSLCP directory (default: from SLTK_DATASETS_ROOT or /vol/research/datasets/multiview/BSLCP)
            split: Split name (BSLCP doesn't have official splits, use 'all')
            feature_root: Path to features directory
            load_poses: Whether to load pose features
            pose_format: Pose format ('wilor', 'smplx', or 'nlf')
            fps: Video FPS (default 25)
            gloss_tiers: List of ELAN tier names to load (default: common gloss tiers)
            signer_filter: Only include specific signers (e.g., ['BF', 'BM'])
        """
        root = root or _get_default_video_root()
        feature_root = feature_root or _get_default_feature_root()

        self.gloss_tiers = gloss_tiers or self.GLOSS_TIERS
        self.signer_filter = signer_filter

        # Store pose format mapping
        self._pose_format = pose_format
        self._feature_subdir = self._get_feature_subdir(pose_format)

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

    def _get_feature_subdir(self, pose_format: str) -> str:
        """Map pose format to feature subdirectory."""
        mapping = {
            "wilor": "wilor",
            "smplx": "smpl",
            "nlf": "smpl",
            "smpl": "smpl",
        }
        return mapping.get(pose_format, pose_format)

    def _load_index(self) -> list[str]:
        """Build index of all BSLCP sessions with EAF annotations."""
        # Find all EAF files
        eaf_files = sorted(self.root.glob("*.eaf"))

        sample_ids = []
        for eaf_path in eaf_files:
            # Extract session ID (e.g., "BF01n" from "BF01n.eaf")
            session_id = eaf_path.stem

            # Apply signer filter if specified
            if self.signer_filter is not None:
                # Signer ID is first 2 characters
                signer = session_id[:2]
                if signer not in self.signer_filter:
                    continue

            # Check that video exists
            video_path = self.root / f"{session_id}.mov"
            if not video_path.exists():
                # Try other extensions
                for ext in [".mp4", ".avi"]:
                    alt_path = self.root / f"{session_id}{ext}"
                    if alt_path.exists():
                        break
                else:
                    # No video found, still include if EAF exists
                    pass

            sample_ids.append(session_id)

        return sample_ids

    def _load_sample(self, sample_id: str) -> Sample:
        """Load a single BSLCP sample."""
        # Find annotation file
        eaf_path = self.root / f"{sample_id}.eaf"
        if not eaf_path.exists():
            raise SampleNotFoundError(sample_id, f"{self.NAME} (missing {eaf_path})")

        # Find video file
        video_path = None
        for ext in [".mov", ".mp4", ".avi"]:
            candidate = self.root / f"{sample_id}{ext}"
            if candidate.exists():
                video_path = candidate
                break

        # Load annotations
        try:
            # Get available tiers
            available_tiers = get_eaf_tiers(eaf_path)

            # Filter to requested gloss tiers that exist
            tiers_to_load = [t for t in self.gloss_tiers if t in available_tiers]

            if tiers_to_load:
                segments = read_eaf(eaf_path, tiers=tiers_to_load)
            else:
                # Load all tiers if no gloss tiers found
                segments = read_eaf(eaf_path)
        except Exception:
            # Return sample with error info
            segments = SegmentList()

        # Extract glosses from segments
        glosses = [s.label for s in segments if s.label]

        # Extract signer ID
        signer_id = self._extract_signer_id(sample_id)

        # Load poses if requested
        poses = None
        if self.load_poses and self.feature_root:
            poses = self._load_poses(sample_id)

        return Sample(
            id=sample_id,
            video_path=video_path,
            poses=poses,
            segments=segments,
            glosses=glosses,
            signer_id=signer_id,
            dataset=self.NAME,
            split=self.split,
            metadata={
                "eaf_path": str(eaf_path),
                "session_type": self._extract_session_type(sample_id),
            },
        )

    def _extract_signer_id(self, sample_id: str) -> str:
        """Extract signer ID from sample ID (e.g., 'BF' from 'BF01n')."""
        # Match pattern: 2 letters at start
        match = re.match(r"^([A-Z]{1,2})", sample_id)
        if match:
            return match.group(1)
        return sample_id[:2]

    def _extract_session_type(self, sample_id: str) -> str:
        """Extract session type from sample ID."""
        # Last character indicates type: n=narrative, c=conversation, etc.
        type_char = sample_id[-1].lower()
        type_map = {
            "n": "narrative",
            "c": "conversation",
            "l": "lexical",
            "i": "interview",
        }
        return type_map.get(type_char, "unknown")

    def _load_poses(self, sample_id: str) -> PoseSequence | None:
        """Load pre-extracted poses for a sample."""
        if self.feature_root is None:
            return None

        # Determine feature file path
        feature_dir = self.feature_root / self._feature_subdir

        # Try common naming patterns
        patterns = [
            f"{sample_id}_{self._pose_format}.h5",
            f"{sample_id}_wilor.h5",
            f"{sample_id}_nlf.h5",
            f"{sample_id}.h5",
        ]

        feature_path = None
        for pattern in patterns:
            candidate = feature_dir / pattern
            if candidate.exists():
                feature_path = candidate
                break

        if feature_path is None:
            return None

        try:
            return PoseSequence.load(
                feature_path,
                format=self._pose_format,
                fps=self.fps,
            )
        except Exception:
            return None

    def get_signers(self) -> list[str]:
        """Get list of unique signer IDs."""
        signers = set()
        for sample_id in self._index:
            signers.add(self._extract_signer_id(sample_id))
        return sorted(signers)

    def filter_by_signer(self, signer: str) -> list[str]:
        """Get sample IDs for a specific signer."""
        return [sid for sid in self._index if self._extract_signer_id(sid) == signer]

    def get_vocabulary(self) -> list[str]:
        """Get unique glosses across all samples."""
        vocab = set()
        for sample in self:
            if sample.glosses:
                vocab.update(sample.glosses)
        return sorted(vocab)
