"""CoNLL-U format reader and core data types."""

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator, Optional


@dataclass
class Token:
    """A single token with all CoNLL-U fields.

    This is the unified token type used throughout the library — for reading
    gold data, for model predictions, and as the public API return type.

    CoNLL-U columns: ID, FORM, LEMMA, UPOS, XPOS, FEATS, HEAD, DEPREL, DEPS, MISC.

    Additional fields (not part of CoNLL-U but used at runtime):
        multiword: If this token is part of a multiword token, a tuple of
            ``(surface_form, (start_id, end_id))`` — e.g.
            ``("Zrobiłbym", (1, 3))`` means this token belongs to MWT
            "Zrobiłbym" spanning word IDs 1, 2, and 3 (inclusive).
        embeddings: Dictionary of named embedding vectors produced by the
            model, e.g. ``{"upostag": [0.1, -0.3, ...]}``.
    """

    id: str
    form: str
    lemma: str = "_"
    upos: str = "_"
    xpos: str = "_"
    feats: dict[str, str] = field(default_factory=dict)
    head: Optional[int] = None
    deprel: str = "_"
    deps: str = "_"
    misc: str = "_"
    multiword: Optional[tuple[str, tuple[int, int]]] = None
    embeddings: dict[str, list[float]] = field(default_factory=dict)

    @classmethod
    def from_line(cls, line: str) -> "Token":
        """Parse a CoNLL-U line into a Token."""
        parts = line.strip().split("\t")
        if len(parts) != 10:
            raise ValueError(f"Expected 10 columns, got {len(parts)}: {line}")

        id_, form, lemma, upos, xpos, feats_str, head_str, deprel, deps, misc = parts

        # Parse features
        feats = {}
        if feats_str != "_":
            for feat in feats_str.split("|"):
                if "=" in feat:
                    key, value = feat.split("=", 1)
                    feats[key] = value

        # Parse head
        head = None
        if head_str != "_":
            try:
                head = int(head_str)
            except ValueError:
                pass

        return cls(
            id=id_,
            form=form,
            lemma=lemma if lemma != "_" else ("_" if "-" in id_ else form),
            upos=upos,
            xpos=xpos,
            feats=feats,
            head=head,
            deprel=deprel,
            deps=deps,
            misc=misc,
        )

    def to_conllu_line(self) -> str:
        """Convert Token to a CoNLL-U format line."""
        feats_str = (
            "|".join(f"{k}={v}" for k, v in sorted(self.feats.items())) if self.feats else "_"
        )
        head_str = str(self.head) if self.head is not None else "_"
        return "\t".join(
            [
                self.id,
                self.form,
                self.lemma,
                self.upos,
                self.xpos,
                feats_str,
                head_str,
                self.deprel,
                self.deps,
                self.misc,
            ]
        )

    # Backward compat alias
    to_line = to_conllu_line

    def to_dict(self) -> dict:
        """Convert Token to a dictionary."""
        d = {
            "id": self.id,
            "form": self.form,
            "lemma": self.lemma,
            "upos": self.upos,
            "xpos": self.xpos,
            "feats": self.feats if self.feats else None,
            "head": self.head,
            "deprel": self.deprel,
        }
        if self.multiword is not None:
            d["multiword"] = {
                "form": self.multiword[0],
                "start": self.multiword[1][0],
                "end": self.multiword[1][1],
            }
        if self.embeddings:
            d["embeddings"] = self.embeddings
        return d

    @property
    def is_multiword(self) -> bool:
        """True if this is a multiword token range (e.g. id "1-2")."""
        return "-" in self.id

    @property
    def is_empty_node(self) -> bool:
        """True if this is an empty node (e.g. id "1.1")."""
        return "." in self.id

    @property
    def word_id(self) -> Optional[int]:
        """Integer word ID, or None for MWT ranges and empty nodes."""
        if self.is_multiword or self.is_empty_node:
            return None
        try:
            return int(self.id)
        except ValueError:
            return None

    def __repr__(self) -> str:
        return f"Token({self.id}, {self.form!r}, upos={self.upos!r}, head={self.head})"


@dataclass
class Sentence:
    """A sentence — a list of tokens with metadata.

    Used for both gold data and model predictions.
    """

    tokens: list[Token]
    sent_id: str = ""
    text: str = ""
    metadata: dict[str, str] = field(default_factory=dict)

    def to_conllu(self) -> str:
        """Convert sentence to CoNLL-U format string.

        If tokens carry ``multiword`` information (but no explicit MWT
        range tokens exist yet), MWT lines are generated automatically.
        """
        lines = []
        if self.sent_id:
            lines.append(f"# sent_id = {self.sent_id}")
        if self.text:
            lines.append(f"# text = {self.text}")
        for key, value in self.metadata.items():
            lines.append(f"# {key} = {value}")

        # Build MWT ranges from token.multiword info
        mwt_ranges: dict[tuple[int, int], str] = {}
        for t in self.tokens:
            if t.multiword is not None:
                surface, span = t.multiword
                mwt_ranges[span] = surface

        emitted_mwt: set[tuple[int, int]] = set()
        for token in self.tokens:
            token_id = int(token.id)
            for (start, end), surface in mwt_ranges.items():
                if token_id == start and (start, end) not in emitted_mwt:
                    mwt_line = f"{start}-{end}\t{surface}" + "\t_" * 8
                    lines.append(mwt_line)
                    emitted_mwt.add((start, end))
            lines.append(token.to_conllu_line())
        lines.append("")
        return "\n".join(lines)

    def to_dict(self) -> dict:
        """Convert Sentence to a dictionary."""
        d = {"sent_id": self.sent_id, "text": self.text}
        d["tokens"] = [t.to_dict() for t in self.tokens]
        return d

    def to_json(self, **kwargs) -> str:
        """Convert Sentence to a JSON string."""
        kwargs.setdefault("ensure_ascii", False)
        return json.dumps(self.to_dict(), **kwargs)

    def __len__(self) -> int:
        """Return number of regular words."""
        return len(self.tokens)

    def __iter__(self) -> Iterator[Token]:
        """Iterate over regular word tokens."""
        return iter(self.tokens)

    def __repr__(self) -> str:
        return f"Sentence(sent_id={self.sent_id!r}, tokens={len(self)}, text={self.text[:50]!r})"


@dataclass
class Document:
    """A collection of sentences — the result of parsing text through a pipeline.

    Each pipeline component stores its output in ``doc.<role>``, where role
    is ``"segmenter"``, ``"parser"``, ``"ner"``, etc.  The top-level
    ``sentences`` always reflects the latest state.

    Example::

        result = nlp("Ala ma kota.")
        result.segmenter   # sentences after segmentation only
        result.parser       # sentences after parsing (with upos, head, …)
        # result.ner        # future: sentences after NER
    """

    sentences: list[Sentence]
    text: str = ""
    treebank_id: str = ""
    language: str = ""
    _annotations: dict[str, list[Sentence]] = field(default_factory=dict, repr=False)

    # ------------------------------------------------------------------
    # Annotation access — doc.segmenter, doc.parser, doc.ner, …
    # ------------------------------------------------------------------

    def __getattr__(self, name: str):
        # Only intercept known pipeline roles stored in _annotations.
        # We must guard against recursion during unpickling / dataclass init.
        if name.startswith("_"):
            raise AttributeError(name)
        try:
            annotations = object.__getattribute__(self, "_annotations")
        except AttributeError:
            raise AttributeError(name)
        if name in annotations:
            return annotations[name]
        raise AttributeError(f"'Document' has no annotation {name!r}. Available: {list(annotations)}")

    def set_annotation(self, role: str, sentences: list["Sentence"]) -> None:
        """Store a snapshot of sentences for a pipeline role.

        Args:
            role: Component role name (e.g. ``"segmenter"``, ``"parser"``).
            sentences: The sentences produced by that component.
        """
        self._annotations[role] = sentences

    @property
    def annotations(self) -> dict[str, list["Sentence"]]:
        """All stored pipeline annotations, keyed by role name."""
        return dict(self._annotations)

    # ------------------------------------------------------------------

    def to_conllu(self) -> str:
        """Convert document to CoNLL-U format string."""
        return "\n".join(sent.to_conllu() for sent in self.sentences)

    def to_dict(self) -> dict:
        """Convert Document to a dictionary."""
        return {"sentences": [s.to_dict() for s in self.sentences]}

    def to_json(self, **kwargs) -> str:
        """Convert Document to a JSON string."""
        kwargs.setdefault("ensure_ascii", False)
        return json.dumps(self.to_dict(), **kwargs)

    def __len__(self) -> int:
        """Return number of sentences."""
        return len(self.sentences)

    def __iter__(self) -> Iterator[Sentence]:
        """Iterate over sentences."""
        return iter(self.sentences)

    def __repr__(self) -> str:
        roles = list(self._annotations.keys())
        roles_str = f", annotations={roles}" if roles else ""
        return f"Document(sentences={len(self.sentences)}{roles_str})"


def _propagate_mwt(tokens: list[Token]) -> None:
    """Set ``multiword`` on sub-word tokens covered by a MWT range token.

    After parsing CoNLL-U lines, MWT range tokens (e.g. ID ``"1-2"``) are
    present in the token list but their sub-words don't know they belong to
    a MWT.  This function sets ``token.multiword = (surface, (start, end))``
    on every sub-word so downstream code can reconstruct MWT lines without
    access to the original range tokens.

    Args:
        tokens: Mutable list of tokens for a single sentence.
    """
    for t in tokens:
        if "-" in t.id:
            start, end = (int(x) for x in t.id.split("-"))
            for sub in tokens:
                if "-" not in sub.id and "." not in sub.id:
                    sub_id = int(sub.id)
                    if start <= sub_id <= end:
                        sub.multiword = (t.form, (start, end))


class CoNLLUReader:
    """Reader for CoNLL-U format files."""

    @staticmethod
    def read(file_path: Path, treebank_id: str = "", language: str = "") -> Document:
        """Read a CoNLL-U file and return a Document.

        Args:
            file_path: Path to the .conllu file.
            treebank_id: Identifier for the treebank (e.g., 'pl_lfg').
            language: Language name (e.g., 'Polish').

        Returns:
            Document containing all sentences.
        """
        sentences = []
        current_tokens = []
        current_metadata = {}
        sent_id = ""
        text = ""

        with open(file_path, encoding="utf-8") as f:
            for line in f:
                line = line.rstrip("\n")

                if not line:
                    # End of sentence
                    if current_tokens:
                        _propagate_mwt(current_tokens)
                        word_tokens = [t for t in current_tokens if not t.is_multiword and not t.is_empty_node]
                        sentences.append(
                            Sentence(
                                sent_id=sent_id or f"sent_{len(sentences)}",
                                text=text,
                                tokens=word_tokens,
                                metadata=current_metadata,
                            )
                        )
                    current_tokens = []
                    current_metadata = {}
                    sent_id = ""
                    text = ""
                elif line.startswith("#"):
                    # Metadata comment
                    if line.startswith("# sent_id"):
                        sent_id = line.split("=", 1)[1].strip()
                    elif line.startswith("# text"):
                        text = line.split("=", 1)[1].strip()
                    else:
                        # Other metadata
                        if "=" in line:
                            key = line[2:].split("=", 1)[0].strip()
                            value = line[2:].split("=", 1)[1].strip()
                            current_metadata[key] = value
                else:
                    # Token line
                    try:
                        token = Token.from_line(line)
                        current_tokens.append(token)
                    except ValueError as e:
                        print(f"Warning: Skipping invalid line in {file_path}: {e}")

        # Handle last sentence if no trailing newline
        if current_tokens:
            _propagate_mwt(current_tokens)
            word_tokens = [t for t in current_tokens if not t.is_multiword and not t.is_empty_node]
            sentences.append(
                Sentence(
                    sent_id=sent_id or f"sent_{len(sentences)}",
                    text=text,
                    tokens=word_tokens,
                    metadata=current_metadata,
                )
            )

        return Document(sentences=sentences, treebank_id=treebank_id, language=language)

    @staticmethod
    def read_multiple(
        file_paths: list[Path], treebank_ids: list[str], language: str
    ) -> list[Document]:
        """Read multiple CoNLL-U files."""
        documents = []
        for file_path, treebank_id in zip(file_paths, treebank_ids):
            doc = CoNLLUReader.read(file_path, treebank_id, language)
            documents.append(doc)
        return documents
