"""Composable NLP processing pipeline."""

from __future__ import annotations

import copy
from typing import Iterable, Iterator, Protocol, runtime_checkable

from .data.conllu_reader import Document, Sentence


@runtime_checkable
class PipelineComponent(Protocol):
    """Any component that can participate in a processing pipeline."""

    def __call__(self, doc: Document) -> Document: ...


class ComponentMixin:
    """Mixin that adds the ``|`` operator and a ``role`` for building pipelines.

    Every component has a ``role`` (e.g. ``"segmenter"``, ``"parser"``,
    ``"ner"``).  When run inside a :class:`Pipeline`, the result of each
    step is stored as ``doc.<role>`` so you can inspect intermediate
    outputs::

        result = nlp("Ala ma kota.")
        result.segmenter   # segmenter output
        result.parser       # parser output (upos, head, deprel, …)
    """

    #: Override in subclasses to set the pipeline role name.
    role: str = "component"

    def __or__(self, other) -> "Pipeline":
        """Chain two components (or a component and a Pipeline) with ``|``."""
        left = self._to_component_list()
        right = other._to_component_list() if hasattr(other, "_to_component_list") else [other]
        return Pipeline(left + right)

    def _to_component_list(self) -> list:
        return [self]


class Pipeline(ComponentMixin):
    """A composable chain of NLP processing components.

    Build pipelines with the ``|`` operator::

        from combo import COMBO
        from combo.segmenters import LamboSegmenter

        nlp = LamboSegmenter("pl") | COMBO.from_pretrained("...")
        result = nlp("Ala ma kota.")

        # Access per-component results:
        result.segmenter   # sentences after segmentation
        result.parser       # sentences after parsing

        # The top-level result always reflects the final state:
        for sentence in result:
            for token in sentence:
                print(token.form, token.upos, token.deprel)
    """

    role = "pipeline"

    def __init__(self, components: list) -> None:
        self.components = list(components)

    def __call__(self, text: str | Document) -> Document:
        """Process text through all pipeline components.

        Each component's output is stored as ``doc.<role>`` before
        passing the document to the next component.

        Args:
            text: Raw text string or a Document object.

        Returns:
            Document with all pipeline components applied.
        """
        if isinstance(text, str):
            doc = Document(sentences=[], text=text)
        else:
            doc = text

        for component in self.components:
            prev_annotations = doc._annotations
            doc = component(doc)
            # Carry forward annotations from previous steps
            for k, v in prev_annotations.items():
                if k not in doc._annotations:
                    doc._annotations[k] = v
            # Store a snapshot of this component's output
            role = getattr(component, "role", type(component).__name__.lower())
            doc.set_annotation(role, copy.deepcopy(doc.sentences))

        return doc

    def pipe(self, texts: Iterable[str], batch_size: int = 32) -> Iterator[Document]:
        """Process multiple texts through the pipeline.

        Args:
            texts: Iterable of text strings.
            batch_size: Batch size (currently processes one at a time).

        Yields:
            Document objects for each input text.
        """
        for text in texts:
            yield self(text)

    def _to_component_list(self) -> list:
        """Flatten this pipeline's components for chaining with ``|``."""
        return list(self.components)

    def __repr__(self) -> str:
        names = [
            getattr(c, "role", type(c).__name__.lower())
            for c in self.components
        ]
        return f"Pipeline([{', '.join(names)}])"
