"""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 ..languages import SplitLevel
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.
        default_split_level: Default split level for segmentation.
            ``"SENTENCE"`` (default) splits turns into sentences.
            ``"TURN"`` treats each turn as a single sentence (no sentence
            splitting within turns). Can be overridden per-call.
        default_split_multiwords: Default for multi-word token splitting.
            When ``True`` (default), MWTs are split into subwords
            (e.g. "can't" → "ca", "n't"). When ``False``, MWTs are kept
            as single tokens. Can be overridden per-call.
    """

    def __init__(
        self,
        language: "str | Language" = "Polish",
        device: str | None = None,
        with_splitter: bool = True,
        model_name: str | None = None,
        default_split_level: "SplitLevel | str" = SplitLevel.SENTENCE,
        default_split_multiwords: bool = True,
    ):
        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
        print(f"  [lambo] loading model: {lambo_name!r}")
        self._lambo = Lambo.get(lambo_name, **kwargs)
        self.language = lang
        self._default_split_level = default_split_level.upper()
        self._default_split_multiwords = default_split_multiwords

    def __call__(
        self,
        doc: Document,
        split_level: "SplitLevel | str | None" = None,
        split_multiwords: bool | None = None,
    ) -> Document:
        """Segment text using LAMBO.

        Args:
            doc: Document with .text set.
            split_level: ``"SENTENCE"`` (default) splits turns into
                sentences. ``"TURN"`` treats each turn as a single
                sentence. Overrides the constructor default.
            split_multiwords: Whether to split multi-word tokens into
                subwords. Overrides the constructor default.

        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)

        level = (split_level.upper() if split_level is not None
                 else self._default_split_level)
        multiwords = (split_multiwords if split_multiwords is not None
                      else self._default_split_multiwords)

        lambo_doc = self._lambo.segment(text)

        sentences = []
        sent_idx = 0

        for turn in lambo_doc.turns:
            if level == "TURN":
                # Treat the entire turn as one sentence — collect all
                # tokens from all LAMBO sentences within this turn.
                all_lambo_tokens = []
                for lambo_sent in turn.sentences:
                    all_lambo_tokens.extend(lambo_sent.tokens)
                if all_lambo_tokens:
                    sent_idx += 1
                    tokens = self._convert_tokens(all_lambo_tokens, multiwords, text)
                    if tokens:
                        sentences.append(
                            Sentence(
                                tokens=tokens,
                                sent_id=f"s{sent_idx}",
                                text=turn.text,
                            )
                        )
            else:
                # Default: each LAMBO sentence becomes a separate sentence.
                for lambo_sent in turn.sentences:
                    sent_idx += 1
                    tokens = self._convert_tokens(lambo_sent.tokens, multiwords, text)
                    if tokens:
                        sentences.append(
                            Sentence(
                                tokens=tokens,
                                sent_id=f"s{sent_idx}",
                                text=lambo_sent.text,
                            )
                        )

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

    @staticmethod
    def _space_after(full_text: str, end: int) -> str:
        """Return ``SpaceAfter=No`` if there is no space at position *end*."""
        if end < len(full_text) and full_text[end] != " ":
            return "SpaceAfter=No"
        return "_"

    @staticmethod
    def _convert_tokens(
        lambo_tokens: list,
        split_multiwords: bool,
        full_text: str,
    ) -> list[Token]:
        """Convert LAMBO tokens to COMBO Token objects.

        Args:
            lambo_tokens: Tokens from LAMBO segmentation.
            split_multiwords: If True, expand MWTs into subword tokens.
                If False, keep MWTs as single tokens.
            full_text: The full document text (used to set SpaceAfter=No).
        """
        tokens: list[Token] = []
        token_idx = 0

        for lambo_token in lambo_tokens:
            misc = LamboSegmenter._space_after(full_text, lambo_token.end)
            if lambo_token.is_multi_word and lambo_token.subwords:
                if split_multiwords:
                    start_idx = token_idx + 1
                    end_idx = start_idx + len(lambo_token.subwords) - 1
                    mwt_info = (lambo_token.text, (start_idx, end_idx))
                    for i, subword in enumerate(lambo_token.subwords):
                        token_idx += 1
                        # SpaceAfter goes on the last subword only
                        subword_misc = misc if i == len(lambo_token.subwords) - 1 else "_"
                        tokens.append(
                            Token(
                                id=str(token_idx),
                                form=subword,
                                multiword=mwt_info,
                                misc=subword_misc,
                            )
                        )
                else:
                    token_idx += 1
                    tokens.append(
                        Token(
                            id=str(token_idx),
                            form=lambo_token.text,
                            misc=misc,
                        )
                    )
            else:
                token_idx += 1
                tokens.append(
                    Token(
                        id=str(token_idx),
                        form=lambo_token.text,
                        misc=misc,
                    )
                )

        return tokens

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