"""Label encoding utilities for UPOS, XPOS, FEATS, and characters."""

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from .conllu_reader import Document, Token


@dataclass
class CombinedLabel:
    """Combined morphological label (UPOS + XPOS + FEATS)."""

    upos: str
    xpos: str
    feats: dict[str, str]

    def to_string(self) -> str:
        """Create canonical string representation for encoding."""
        feats_str = "|".join(f"{k}={v}" for k, v in sorted(self.feats.items()))
        return f"{self.upos}||{self.xpos}||{feats_str}"

    @classmethod
    def from_string(cls, label_str: str) -> "CombinedLabel":
        """Parse a string representation back to CombinedLabel."""
        parts = label_str.split("||")
        if len(parts) != 3:
            return cls(upos="_", xpos="_", feats={})

        upos, xpos, feats_str = parts
        feats = {}
        if feats_str:
            for feat in feats_str.split("|"):
                if "=" in feat:
                    key, value = feat.split("=", 1)
                    feats[key] = value
        return cls(upos=upos, xpos=xpos, feats=feats)

    @classmethod
    def from_token(cls, token: Token) -> "CombinedLabel":
        """Create CombinedLabel from a Token."""
        return cls(upos=token.upos, xpos=token.xpos, feats=token.feats)


class CombinedLabelEncoder:
    """Encoder for combined UPOS+XPOS+FEATS labels."""

    PAD_TOKEN = "<PAD>"
    UNK_TOKEN = "<UNK>"

    def __init__(self):
        self.label_to_id: dict[str, int] = {self.PAD_TOKEN: 0, self.UNK_TOKEN: 1}
        self.id_to_label: dict[int, str] = {0: self.PAD_TOKEN, 1: self.UNK_TOKEN}
        self._frozen = False

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

    def fit(self, documents: list[Document]) -> "CombinedLabelEncoder":
        """Build vocabulary from training data.

        Args:
            documents: List of training documents.

        Returns:
            Self for chaining.
        """
        for doc in documents:
            for sentence in doc.sentences:
                for token in sentence.tokens:
                    label = CombinedLabel.from_token(token)
                    label_str = label.to_string()
                    if label_str not in self.label_to_id:
                        idx = len(self.label_to_id)
                        self.label_to_id[label_str] = idx
                        self.id_to_label[idx] = label_str
        self._frozen = True
        return self

    def encode(self, label: CombinedLabel) -> int:
        """Encode a CombinedLabel to integer ID."""
        label_str = label.to_string()
        return self.label_to_id.get(label_str, self.label_to_id[self.UNK_TOKEN])

    def decode(self, label_id: int) -> str:
        """Decode an integer ID to label string."""
        return self.id_to_label.get(label_id, self.UNK_TOKEN)

    def decode_to_components(self, label_id: int) -> tuple[str, str, dict[str, str]]:
        """Decode back to UPOS, XPOS, FEATS components."""
        label_str = self.decode(label_id)
        if label_str in (self.PAD_TOKEN, self.UNK_TOKEN):
            return "_", "_", {}
        label = CombinedLabel.from_string(label_str)
        return label.upos, label.xpos, label.feats

    def save(self, path: Path) -> None:
        """Save encoder to JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump({"label_to_id": self.label_to_id}, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> "CombinedLabelEncoder":
        """Load encoder from JSON file."""
        encoder = cls()
        with open(path) as f:
            data = json.load(f)
        encoder.label_to_id = data["label_to_id"]
        encoder.id_to_label = {int(v): k for k, v in encoder.label_to_id.items()}
        encoder._frozen = True
        return encoder


class LabelEncoder:
    """Simple label encoder for single-value labels (e.g., DEPREL)."""

    PAD_TOKEN = "<PAD>"
    UNK_TOKEN = "<UNK>"

    def __init__(self):
        self.label_to_id: dict[str, int] = {self.PAD_TOKEN: 0, self.UNK_TOKEN: 1}
        self.id_to_label: dict[int, str] = {0: self.PAD_TOKEN, 1: self.UNK_TOKEN}
        self._frozen = False

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

    def fit(self, labels: list[str]) -> "LabelEncoder":
        """Build vocabulary from labels.

        Args:
            labels: List of label strings.

        Returns:
            Self for chaining.
        """
        for label in labels:
            if label not in self.label_to_id:
                idx = len(self.label_to_id)
                self.label_to_id[label] = idx
                self.id_to_label[idx] = label
        self._frozen = True
        return self

    def fit_documents(self, documents: list[Document], field: str = "deprel") -> "LabelEncoder":
        """Build vocabulary from documents.

        Args:
            documents: List of documents.
            field: Token field to extract ('deprel', 'upos', etc.).

        Returns:
            Self for chaining.
        """
        labels = []
        for doc in documents:
            for sentence in doc.sentences:
                for token in sentence.tokens:
                    labels.append(getattr(token, field))
        return self.fit(labels)

    def encode(self, label: str) -> int:
        """Encode a label to integer ID."""
        return self.label_to_id.get(label, self.label_to_id[self.UNK_TOKEN])

    def decode(self, label_id: int) -> str:
        """Decode an integer ID to label string."""
        return self.id_to_label.get(label_id, self.UNK_TOKEN)

    def save(self, path: Path) -> None:
        """Save encoder to JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump({"label_to_id": self.label_to_id}, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> "LabelEncoder":
        """Load encoder from JSON file."""
        encoder = cls()
        with open(path) as f:
            data = json.load(f)
        encoder.label_to_id = data["label_to_id"]
        encoder.id_to_label = {int(v): k for k, v in encoder.label_to_id.items()}
        encoder._frozen = True
        return encoder


class MorphoEncoder:
    """Decomposed morphological encoder with separate vocabularies for UPOS, XPOS, and each feature category."""

    def __init__(self):
        self.upos_encoder = LabelEncoder()
        self.xpos_encoder = LabelEncoder()
        self.feat_categories: list[str] = []  # sorted list of all feat category names
        self.feat_encoders: dict[str, LabelEncoder] = {}
        self.valid_feat_vectors: dict[int, set[tuple[int, ...]]] = {}  # upos_id -> valid feat combos
        self._frozen = False

    @property
    def num_upos(self) -> int:
        return len(self.upos_encoder)

    @property
    def num_xpos(self) -> int:
        return len(self.xpos_encoder)

    @property
    def feat_vocab_sizes(self) -> dict[str, int]:
        return {cat: len(self.feat_encoders[cat]) for cat in self.feat_categories}

    @property
    def num_feat_categories(self) -> int:
        return len(self.feat_categories)

    def fit(self, documents: list[Document]) -> "MorphoEncoder":
        """Build vocabularies from training data."""
        upos_labels = []
        xpos_labels = []
        feat_values: dict[str, list[str]] = {}  # cat -> list of values

        for doc in documents:
            for sentence in doc.sentences:
                for token in sentence.tokens:
                    upos_labels.append(token.upos)
                    xpos_labels.append(token.xpos)
                    for key, value in token.feats.items():
                        if key not in feat_values:
                            feat_values[key] = []
                        feat_values[key].append(value)

        self.upos_encoder.fit(upos_labels)
        self.xpos_encoder.fit(xpos_labels)

        self.feat_categories = sorted(feat_values.keys())
        for cat in self.feat_categories:
            encoder = LabelEncoder()
            encoder.fit(feat_values[cat])
            self.feat_encoders[cat] = encoder

        # Build valid (UPOS -> feat_vector) combinations for constraint filtering
        self._build_valid_combinations(documents)

        self._frozen = True
        return self

    def _build_valid_combinations(self, documents: list[Document]) -> None:
        """Build lookup of valid (UPOS -> feat_vector) combinations from training data."""
        self.valid_feat_vectors = {}
        for doc in documents:
            for sentence in doc.sentences:
                for token in sentence.tokens:
                    upos_id = self.encode_upos(token.upos)
                    feat_ids = tuple(self.encode_feats(token.feats))
                    if upos_id not in self.valid_feat_vectors:
                        self.valid_feat_vectors[upos_id] = set()
                    self.valid_feat_vectors[upos_id].add(feat_ids)

    def constrain_feats(self, upos_id: int, feat_ids: list[int]) -> list[int]:
        """If (upos_id, feat_ids) is not a valid combination, find closest valid one.

        Args:
            upos_id: Predicted UPOS ID.
            feat_ids: Predicted feature IDs (one per category).

        Returns:
            Possibly corrected feature IDs.
        """
        feat_tuple = tuple(feat_ids)
        valid_set = self.valid_feat_vectors.get(upos_id, set())

        if not valid_set or feat_tuple in valid_set:
            return feat_ids

        # Find closest by Hamming distance
        best_dist = len(feat_ids) + 1
        best_vec = feat_tuple
        for valid_vec in valid_set:
            dist = sum(1 for a, b in zip(feat_tuple, valid_vec) if a != b)
            if dist < best_dist:
                best_dist = dist
                best_vec = valid_vec
                if dist == 0:
                    break
        return list(best_vec)

    def encode_upos(self, upos: str) -> int:
        return self.upos_encoder.encode(upos)

    def encode_xpos(self, xpos: str) -> int:
        return self.xpos_encoder.encode(xpos)

    def encode_feats(self, feats: dict[str, str]) -> list[int]:
        """Encode feats dict to fixed-length vector (one int per category, 0=not applicable)."""
        result = []
        for cat in self.feat_categories:
            if cat in feats:
                result.append(self.feat_encoders[cat].encode(feats[cat]))
            else:
                result.append(0)  # PAD = not applicable
        return result

    def decode_upos(self, label_id: int) -> str:
        return self.upos_encoder.decode(label_id)

    def decode_xpos(self, label_id: int) -> str:
        return self.xpos_encoder.decode(label_id)

    def decode_feats(self, ids: list[int]) -> dict[str, str]:
        """Decode fixed-length vector back to feats dict."""
        feats = {}
        for i, cat in enumerate(self.feat_categories):
            if ids[i] != 0:  # 0 = PAD = not applicable
                value = self.feat_encoders[cat].decode(ids[i])
                if value not in (LabelEncoder.PAD_TOKEN, LabelEncoder.UNK_TOKEN):
                    feats[cat] = value
        return feats

    def save(self, path: Path) -> None:
        """Save encoder to JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "upos": {"label_to_id": self.upos_encoder.label_to_id},
            "xpos": {"label_to_id": self.xpos_encoder.label_to_id},
            "feat_categories": self.feat_categories,
            "feats": {
                cat: {"label_to_id": self.feat_encoders[cat].label_to_id}
                for cat in self.feat_categories
            },
            "valid_feat_vectors": {
                str(upos_id): [list(v) for v in vectors]
                for upos_id, vectors in self.valid_feat_vectors.items()
            },
        }
        with open(path, "w") as f:
            json.dump(data, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> "MorphoEncoder":
        """Load encoder from JSON file."""
        encoder = cls()
        with open(path) as f:
            data = json.load(f)

        # Load UPOS encoder
        encoder.upos_encoder.label_to_id = data["upos"]["label_to_id"]
        encoder.upos_encoder.id_to_label = {
            int(v): k for k, v in encoder.upos_encoder.label_to_id.items()
        }
        encoder.upos_encoder._frozen = True

        # Load XPOS encoder
        encoder.xpos_encoder.label_to_id = data["xpos"]["label_to_id"]
        encoder.xpos_encoder.id_to_label = {
            int(v): k for k, v in encoder.xpos_encoder.label_to_id.items()
        }
        encoder.xpos_encoder._frozen = True

        # Load feat encoders
        encoder.feat_categories = data["feat_categories"]
        for cat in encoder.feat_categories:
            feat_enc = LabelEncoder()
            feat_enc.label_to_id = data["feats"][cat]["label_to_id"]
            feat_enc.id_to_label = {
                int(v): k for k, v in feat_enc.label_to_id.items()
            }
            feat_enc._frozen = True
            encoder.feat_encoders[cat] = feat_enc

        # Load valid feat vector combinations
        if "valid_feat_vectors" in data:
            encoder.valid_feat_vectors = {
                int(upos_id): {tuple(v) for v in vectors}
                for upos_id, vectors in data["valid_feat_vectors"].items()
            }

        encoder._frozen = True
        return encoder


class CharVocab:
    """Character vocabulary for lemmatization."""

    PAD_TOKEN = "<PAD>"
    UNK_TOKEN = "<UNK>"
    SOS_TOKEN = "<SOS>"
    EOS_TOKEN = "<EOS>"

    def __init__(self, max_len: int = 50):
        self.char_to_id: dict[str, int] = {
            self.PAD_TOKEN: 0,
            self.UNK_TOKEN: 1,
            self.SOS_TOKEN: 2,
            self.EOS_TOKEN: 3,
        }
        self.id_to_char: dict[int, str] = {
            0: self.PAD_TOKEN,
            1: self.UNK_TOKEN,
            2: self.SOS_TOKEN,
            3: self.EOS_TOKEN,
        }
        self.max_len = max_len
        self._frozen = False

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

    @property
    def pad_id(self) -> int:
        return self.char_to_id[self.PAD_TOKEN]

    @property
    def unk_id(self) -> int:
        return self.char_to_id[self.UNK_TOKEN]

    @property
    def sos_id(self) -> int:
        return self.char_to_id[self.SOS_TOKEN]

    @property
    def eos_id(self) -> int:
        return self.char_to_id[self.EOS_TOKEN]

    def fit(self, documents: list[Document]) -> "CharVocab":
        """Build vocabulary from documents.

        Args:
            documents: List of training documents.

        Returns:
            Self for chaining.
        """
        for doc in documents:
            for sentence in doc.sentences:
                for token in sentence.tokens:
                    for char in token.form:
                        if char not in self.char_to_id:
                            idx = len(self.char_to_id)
                            self.char_to_id[char] = idx
                            self.id_to_char[idx] = char
                    for char in token.lemma:
                        if char not in self.char_to_id:
                            idx = len(self.char_to_id)
                            self.char_to_id[char] = idx
                            self.id_to_char[idx] = char
        self._frozen = True
        return self

    def encode(self, text: str, add_special: bool = False) -> list[int]:
        """Encode text to list of character IDs.

        Args:
            text: Input text.
            add_special: Whether to add SOS/EOS tokens.

        Returns:
            List of character IDs.
        """
        ids = []
        if add_special:
            ids.append(self.sos_id)

        for char in text[: self.max_len - (2 if add_special else 0)]:
            ids.append(self.char_to_id.get(char, self.unk_id))

        if add_special:
            ids.append(self.eos_id)

        return ids

    def decode(self, ids: list[int], remove_special: bool = True) -> str:
        """Decode list of character IDs to text.

        Args:
            ids: List of character IDs.
            remove_special: Whether to remove special tokens.

        Returns:
            Decoded text.
        """
        chars = []
        for id_ in ids:
            char = self.id_to_char.get(id_, self.UNK_TOKEN)
            if remove_special:
                if char in (self.PAD_TOKEN, self.UNK_TOKEN, self.SOS_TOKEN):
                    continue
                if char == self.EOS_TOKEN:
                    break
            chars.append(char)
        return "".join(chars)

    def save(self, path: Path) -> None:
        """Save vocabulary to JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump({"char_to_id": self.char_to_id, "max_len": self.max_len}, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> "CharVocab":
        """Load vocabulary from JSON file."""
        with open(path) as f:
            data = json.load(f)
        vocab = cls(max_len=data["max_len"])
        vocab.char_to_id = data["char_to_id"]
        vocab.id_to_char = {int(v): k for k, v in vocab.char_to_id.items()}
        vocab._frozen = True
        return vocab
