"""LAMBO segmenter integration.

LAMBO (Layered Approach to Multi-level BOundary identification) is a neural
segmentation tool supporting 67 languages. Install it with::

    pip install --index-url https://pypi.clarin-pl.eu/ lambo

Usage::

    from combo.segmenters import LamboSegmenter
    segmenter = LamboSegmenter("pl")
"""

from ..data.conllu_reader import Document, Sentence, Token
from .base import BaseSegmenter


class LamboSegmenter(BaseSegmenter):
    """Segmenter using the LAMBO library.

    Args:
        language: Language enum member or language name string
            (e.g. Language.POLISH, "Polish").
        device: PyTorch device for the segmenter model.
        with_splitter: Whether to enable subword splitting.
        model_name: Full LAMBO model name (e.g. "LAMBO_213-UD_Polish-PDB").
            When provided, this is passed directly to ``Lambo.get()``
            instead of the language name.
    """

    def __init__(
        self,
        language: "str | Language" = "Polish",
        device: str | None = None,
        with_splitter: bool = True,
        model_name: str | None = None,
    ):
        from ..languages import Language, resolve_language

        lang = resolve_language(language)

        try:
            from lambo.segmenter.lambo import Lambo
        except ImportError:
            raise ImportError(
                "LAMBO is not installed. Install it with:\n"
                "  pip install --index-url https://pypi.clarin-pl.eu/ lambo\n"
                "Or use COMBO.from_pretrained() with pre-tokenized input (list of strings)."
            )

        kwargs = {}
        if device is not None:
            kwargs["device"] = device
        if not with_splitter:
            kwargs["with_splitter"] = False

        lambo_name = model_name if model_name is not None else lang.value
        self._lambo = Lambo.get(lambo_name, **kwargs)
        self.language = lang

    def __call__(self, doc: Document) -> Document:
        """Segment text using LAMBO.

        Args:
            doc: Document with .text set.

        Returns:
            Document with sentences and tokens from LAMBO segmentation.
        """
        if doc.sentences:
            return doc

        text = doc.text.strip()
        if not text:
            return Document(sentences=[], text=doc.text)

        lambo_doc = self._lambo.segment(text)

        sentences = []
        sent_idx = 0
        for turn in lambo_doc.turns:
            for lambo_sent in turn.sentences:
                sent_idx += 1
                tokens = []
                token_idx = 0
                for lambo_token in lambo_sent.tokens:
                    token_idx += 1
                    tokens.append(
                        Token(
                            id=str(token_idx),
                            form=lambo_token.text,
                        )
                    )
                if tokens:
                    sent_text = " ".join(t.form for t in tokens)
                    sentences.append(
                        Sentence(
                            tokens=tokens,
                            sent_id=f"s{sent_idx}",
                            text=sent_text,
                        )
                    )

        return Document(sentences=sentences, text=doc.text)

    def __repr__(self) -> str:
        return f"LamboSegmenter({self.language.value!r})"
