"""
Vocabulary management for sign language glosses.

Provides utilities for building, saving, and loading gloss vocabularies,
as well as gloss-to-ID conversion for model training.
"""

import json
from collections import Counter
from collections.abc import Iterator
from pathlib import Path

from sltk.data.core import Sample, SegmentList


class Vocabulary:
    """
    Gloss vocabulary for sign language datasets.

    Handles mapping between gloss strings and integer IDs,
    with support for special tokens (PAD, UNK, BOS, EOS).

    Example:
        >>> vocab = Vocabulary.from_segments(segments)
        >>> ids = vocab.encode(["HELLO", "WORLD"])
        >>> glosses = vocab.decode([5, 10])
    """

    # Special token defaults
    PAD_TOKEN = "<pad>"
    UNK_TOKEN = "<unk>"
    BOS_TOKEN = "<bos>"
    EOS_TOKEN = "<eos>"

    def __init__(
        self,
        glosses: list[str] | None = None,
        min_count: int = 1,
        add_special_tokens: bool = True,
    ):
        """
        Initialize vocabulary.

        Args:
            glosses: Optional list of gloss strings to build vocab from
            min_count: Minimum count to include a gloss
            add_special_tokens: Whether to add PAD, UNK, BOS, EOS tokens
        """
        self._gloss_to_id: dict[str, int] = {}
        self._id_to_gloss: dict[int, str] = {}
        self._counts: dict[str, int] = {}
        self._min_count = min_count

        # Add special tokens first
        if add_special_tokens:
            for token in [self.PAD_TOKEN, self.UNK_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN]:
                self._add_gloss(token)

        if glosses is not None:
            self.build(glosses, min_count)

    def _add_gloss(self, gloss: str) -> int:
        """Add a single gloss to vocabulary."""
        if gloss not in self._gloss_to_id:
            idx = len(self._gloss_to_id)
            self._gloss_to_id[gloss] = idx
            self._id_to_gloss[idx] = gloss
        return self._gloss_to_id[gloss]

    def build(self, glosses: list[str], min_count: int = 1) -> "Vocabulary":
        """
        Build vocabulary from list of glosses.

        Args:
            glosses: List of gloss strings
            min_count: Minimum count to include

        Returns:
            self for chaining
        """
        # Count occurrences
        counts = Counter(glosses)
        self._counts.update(counts)

        # Add glosses meeting threshold
        for gloss, count in sorted(counts.items()):
            if count >= min_count:
                self._add_gloss(gloss)

        return self

    @classmethod
    def from_segments(
        cls,
        segments: SegmentList,
        min_count: int = 1,
        **kwargs,
    ) -> "Vocabulary":
        """
        Build vocabulary from segment annotations.

        Args:
            segments: SegmentList with labeled segments
            min_count: Minimum count to include

        Returns:
            New Vocabulary instance
        """
        glosses = [s.label for s in segments if s.label is not None]
        return cls(glosses, min_count=min_count, **kwargs)

    @classmethod
    def from_samples(
        cls,
        samples: Iterator[Sample],
        min_count: int = 1,
        use_segments: bool = True,
        **kwargs,
    ) -> "Vocabulary":
        """
        Build vocabulary from dataset samples.

        Args:
            samples: Iterator of Sample objects
            min_count: Minimum count to include
            use_segments: If True, extract from segments; else use glosses list

        Returns:
            New Vocabulary instance
        """
        all_glosses: list[str] = []
        for sample in samples:
            if use_segments and sample.segments:
                all_glosses.extend(s.label for s in sample.segments if s.label)
            elif sample.glosses:
                all_glosses.extend(sample.glosses)

        return cls(all_glosses, min_count=min_count, **kwargs)

    def encode(
        self,
        glosses: str | list[str],
        add_bos: bool = False,
        add_eos: bool = False,
    ) -> list[int]:
        """
        Convert gloss(es) to ID(s).

        Args:
            glosses: Single gloss or list of glosses
            add_bos: Add beginning-of-sequence token
            add_eos: Add end-of-sequence token

        Returns:
            List of IDs
        """
        if isinstance(glosses, str):
            glosses = [glosses]

        unk_id = self._gloss_to_id.get(self.UNK_TOKEN, 0)
        ids = [self._gloss_to_id.get(g, unk_id) for g in glosses]

        if add_bos:
            ids = [self._gloss_to_id[self.BOS_TOKEN]] + ids
        if add_eos:
            ids = ids + [self._gloss_to_id[self.EOS_TOKEN]]

        return ids

    def decode(
        self,
        ids: list[int],
        skip_special: bool = True,
    ) -> list[str]:
        """
        Convert ID(s) to gloss(es).

        Args:
            ids: List of IDs
            skip_special: Skip special tokens in output

        Returns:
            List of glosses
        """
        special = {self.PAD_TOKEN, self.UNK_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN}
        glosses = []

        for idx in ids:
            gloss = self._id_to_gloss.get(idx, self.UNK_TOKEN)
            if skip_special and gloss in special:
                continue
            glosses.append(gloss)

        return glosses

    def __len__(self) -> int:
        return len(self._gloss_to_id)

    def __contains__(self, gloss: str) -> bool:
        return gloss in self._gloss_to_id

    def __getitem__(self, key: str | int) -> int | str:
        if isinstance(key, str):
            return self._gloss_to_id[key]
        return self._id_to_gloss[key]

    @property
    def size(self) -> int:
        """Vocabulary size including special tokens."""
        return len(self)

    @property
    def glosses(self) -> list[str]:
        """List of all glosses in vocabulary order."""
        return [self._id_to_gloss[i] for i in range(len(self))]

    @property
    def pad_id(self) -> int:
        return self._gloss_to_id.get(self.PAD_TOKEN, 0)

    @property
    def unk_id(self) -> int:
        return self._gloss_to_id.get(self.UNK_TOKEN, 1)

    @property
    def bos_id(self) -> int:
        return self._gloss_to_id.get(self.BOS_TOKEN, 2)

    @property
    def eos_id(self) -> int:
        return self._gloss_to_id.get(self.EOS_TOKEN, 3)

    def get_count(self, gloss: str) -> int:
        """Get occurrence count for a gloss."""
        return self._counts.get(gloss, 0)

    def most_common(self, n: int = 10) -> list[tuple]:
        """Get n most common glosses."""
        sorted_counts = sorted(self._counts.items(), key=lambda x: x[1], reverse=True)
        return sorted_counts[:n]

    def save(self, path: str | Path) -> Path:
        """
        Save vocabulary to JSON file.

        Args:
            path: Output path

        Returns:
            Path to saved file
        """
        path = Path(path)
        data = {
            "glosses": self.glosses,
            "counts": self._counts,
            "min_count": self._min_count,
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        return path

    @classmethod
    def load(cls, path: str | Path) -> "Vocabulary":
        """
        Load vocabulary from JSON file.

        Args:
            path: Path to vocabulary file

        Returns:
            Loaded Vocabulary instance
        """
        with open(path, encoding="utf-8") as f:
            data = json.load(f)

        vocab = cls(add_special_tokens=False)
        for gloss in data["glosses"]:
            vocab._add_gloss(gloss)
        vocab._counts = data.get("counts", {})
        vocab._min_count = data.get("min_count", 1)

        return vocab

    def __repr__(self) -> str:
        return f"Vocabulary(size={len(self)}, min_count={self._min_count})"


def normalize_gloss(gloss: str, lowercase: bool = True) -> str:
    """
    Normalize a gloss string.

    Standard normalization:
    - Strip whitespace
    - Optionally lowercase
    - Remove extra spaces

    Args:
        gloss: Input gloss string
        lowercase: Whether to convert to lowercase

    Returns:
        Normalized gloss
    """
    gloss = gloss.strip()
    gloss = " ".join(gloss.split())  # Normalize whitespace

    if lowercase:
        gloss = gloss.lower()

    return gloss


def parse_compound_gloss(gloss: str, separator: str = "+") -> list[str]:
    """
    Parse a compound gloss into components.

    Example:
        >>> parse_compound_gloss("HELLO+WORLD")
        ['HELLO', 'WORLD']

    Args:
        gloss: Compound gloss string
        separator: Separator character

    Returns:
        List of component glosses
    """
    return [g.strip() for g in gloss.split(separator) if g.strip()]


def is_pointing_sign(gloss: str) -> bool:
    """
    Check if a gloss is a pointing/pronominal sign.

    Common patterns: PT:PRO1SG, IX, POINT, etc.

    Args:
        gloss: Gloss string

    Returns:
        True if pointing sign
    """
    gloss_upper = gloss.upper()
    pointing_patterns = ["PT:", "IX", "POINT", "PRO", "POSS"]
    return any(gloss_upper.startswith(p) for p in pointing_patterns)


def is_fingerspelling(gloss: str) -> bool:
    """
    Check if a gloss represents fingerspelling.

    Common patterns: FS:word, #word, fs-word

    Args:
        gloss: Gloss string

    Returns:
        True if fingerspelling
    """
    gloss_upper = gloss.upper()
    fs_patterns = ["FS:", "FS-", "#"]
    return any(gloss_upper.startswith(p) for p in fs_patterns)


def gloss_category(gloss: str) -> str:
    """
    Categorize a gloss by type.

    Categories:
    - "pointing": Pronominal/pointing signs
    - "fingerspelling": Fingerspelled words
    - "compound": Compound signs (contains +)
    - "lexical": Regular lexical signs

    Args:
        gloss: Gloss string

    Returns:
        Category string
    """
    if is_pointing_sign(gloss):
        return "pointing"
    if is_fingerspelling(gloss):
        return "fingerspelling"
    if "+" in gloss:
        return "compound"
    return "lexical"
