"""COMBO — the main user-facing API.

Example::

    from combo import COMBO, Language

    # Quick start with language enum — auto-segments raw text
    nlp = COMBO(Language.POLISH)
    result = nlp("Ala ma kota.")

    # Pre-tokenized input (no segmenter needed)
    nlp = COMBO.from_pretrained("clarin-pl/combo-polish-herbert-base")
    result = nlp(["Ala", "ma", "kota", "."], tokenized=True)
    result = nlp([["Ala", "ma", "kota", "."], ["Pies", "je", "."]], tokenized=True)

    # Access results
    for sentence in result:
        for token in sentence:
            print(token.form, token.upos, token.head, token.deprel, token.lemma)
"""

from pathlib import Path
from typing import Optional

import torch

from .data.conllu_reader import Document, Sentence, Token
from .languages import SplitLevel
from .pipeline import ComponentMixin, Pipeline


class COMBO(ComponentMixin):
    """COMBO Universal Dependencies parser.

    Works both as a standalone parser and as a component in a Pipeline.
    Supports the ``|`` operator::

        nlp = LamboSegmenter(Language.POLISH) | COMBO(Language.POLISH)
    """

    role = "parser"

    def __init__(
        self,
        language: "str | Language",
        device: Optional[str] = None,
        batch_size: int = 1,
    ):
        """Initialize COMBO with a language name or enum.

        Args:
            language: Language enum member (e.g. Language.POLISH) or
                language name string (e.g. "Polish").
            device: Device string ("cuda", "cpu", "mps", "auto").
                Auto-detected if None.
            batch_size: Number of sentences to process in parallel during
                inference. Higher values improve GPU utilization but use more
                memory. Default is 1 (sequential, backward-compatible).

        Example::

            from combo import COMBO, Language
            nlp = COMBO(Language.POLISH)
            nlp = COMBO("Polish", batch_size=32)
            result = nlp("Ala ma kota.")

        Raises:
            UnsupportedLanguageError: If the language string is not recognized.
            LanguageNotTrainedError: If the language is recognized but has no
                trained COMBO model available yet.
        """
        from .languages import Language, resolve_language

        self._segmenter = None
        self._batch_size = batch_size

        lang = resolve_language(language)
        from .hub import resolve_model, resolve_lambo_model
        model_id = resolve_model(lang.value)
        instance = COMBO.from_pretrained(model_id, device=device, batch_size=batch_size)
        self._predictor = instance._predictor
        self._batch_size = instance._batch_size

        from .segmenters import LamboSegmenter
        lambo_model = resolve_lambo_model(model_id)
        self._segmenter = LamboSegmenter(lang, model_name=lambo_model)

    @classmethod
    def from_predictor(
        cls,
        predictor,
        batch_size: int = 1,
    ) -> "COMBO":
        """Create a COMBO instance from an existing Predictor.

        Use this when you have a custom or locally trained model and want full
        control over how the predictor is loaded (device, sequence length, etc.).
        No segmenter is attached — raw text input is not available; use
        pre-tokenized input (``tokenized=True``) instead.

        Args:
            predictor: A :class:`~combo.inference.Predictor` instance.
            batch_size: Number of sentences per forward pass. Default is 1.

        Returns:
            COMBO instance ready for parsing.

        Example::

            from pathlib import Path
            import torch
            from combo import COMBO
            from combo.inference.predictor import Predictor
            from combo.config import ModelConfig

            config = ModelConfig(...)
            predictor = Predictor.from_checkpoint(
                checkpoint_path=Path("path/to/checkpoint.pt"),
                encoders_dir=Path("path/to/encoders/"),
                config=config,
                device=torch.device("cuda"),
            )
            nlp = COMBO.from_predictor(predictor)
            result = nlp(["Ala", "ma", "kota", "."], tokenized=True)
        """
        instance = cls.__new__(cls)
        instance._predictor = predictor
        instance._segmenter = None
        instance._batch_size = batch_size
        return instance

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        device: Optional[str] = None,
        batch_size: int = 1,
    ) -> "COMBO":
        """Load a COMBO model from HuggingFace Hub or a local path.

        Args:
            model_name_or_path: Either a HuggingFace model ID
                (e.g. "clarin-pl/combo-polish-herbert-base") or a local directory
                path containing pytorch_model.bin, config.json, and encoders/.
            device: Device string. Auto-detected if None.

        Returns:
            COMBO instance ready for parsing.

        Example::

            from combo import COMBO
            nlp = COMBO.from_pretrained("clarin-pl/combo-polish-herbert-base")
            sent = nlp.parse(["Ala", "ma", "kota"])
        """
        from .config import load_config
        from .inference import Predictor

        # Resolve device
        if device is None:
            if torch.cuda.is_available():
                torch_device = torch.device("cuda")
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                torch_device = torch.device("mps")
            else:
                torch_device = torch.device("cpu")
        else:
            torch_device = torch.device(device)

        # Determine local path
        local_path = Path(model_name_or_path)
        if not local_path.exists():
            # Download from HuggingFace Hub
            from .hub import download_model
            local_path = download_model(model_name_or_path)

        # Load config
        config_path = local_path / "config.json"
        if not config_path.exists():
            raise FileNotFoundError(
                f"config.json not found in {local_path}. "
                f"Expected model directory to contain: config.json, pytorch_model.bin, encoders/"
            )
        config = load_config(config_path)

        # Load weights
        checkpoint_path = local_path / "pytorch_model.bin"
        if not checkpoint_path.exists():
            raise FileNotFoundError(
                f"pytorch_model.bin not found in {local_path}."
            )

        encoders_dir = local_path / "encoders"

        predictor = Predictor.from_checkpoint(
            checkpoint_path=checkpoint_path,
            encoders_dir=encoders_dir,
            config=config.model,
            device=torch_device,
        )

        instance = cls.from_predictor(predictor, batch_size=batch_size)

        # Try to set up a default segmenter from the registry
        from .hub import resolve_lambo_model, resolve_language_for_model
        lambo_model = resolve_lambo_model(model_name_or_path)
        lang_name = resolve_language_for_model(model_name_or_path)
        if lambo_model and lang_name:
            try:
                from .segmenters import LamboSegmenter
                instance._segmenter = LamboSegmenter(lang_name, model_name=lambo_model)
            except Exception:
                pass  # LAMBO not installed or model not available

        return instance

    def __call__(
        self,
        input: "str | list[str] | list[list[str]] | Document",
        tokenized: bool = False,
        split_level: "SplitLevel | str | None" = None,
        split_multiwords: bool | None = None,
    ) -> Document:
        """Parse input.

        Args:
            input:
                - ``str``: Raw text (requires segmenter, ``tokenized`` must be False).
                - ``list[str]``: When ``tokenized=False`` — list of raw sentences
                  (each segmented separately). When ``tokenized=True`` — single
                  pre-tokenized sentence.
                - ``list[list[str]]``: Multiple pre-tokenized sentences
                  (``tokenized=True``).
                - ``Document``: From a previous pipeline component.
            tokenized: If True, treat list input as pre-tokenized (skip
                segmentation). Default False.
            split_level: Controls sentence splitting granularity.
                ``"SENTENCE"`` (default) — split turns into sentences.
                ``"TURN"`` — treat each turn as one sentence (no sentence
                splitting within turns, useful for speech/dialogue).
                Only used when segmenter is active (raw text input).
            split_multiwords: Whether to split multi-word tokens into
                subwords (e.g. "can't" → "ca", "n't"). Default True.
                Only used when segmenter is active (raw text input).

        Returns:
            Document with parsed sentences.
        """
        seg_kwargs = {}
        if split_level is not None:
            seg_kwargs["split_level"] = split_level
        if split_multiwords is not None:
            seg_kwargs["split_multiwords"] = split_multiwords

        if isinstance(input, Document):
            return self._predict_doc(input)

        if isinstance(input, str):
            if tokenized:
                raise TypeError("tokenized=True requires list input, not str.")
            if self._segmenter is None:
                raise TypeError(
                    "Raw text input requires a segmenter. Initialize "
                    "COMBO with a language (e.g. COMBO('Polish'))."
                )
            doc = self._segmenter(Document(sentences=[], text=input), **seg_kwargs)
            return self._predict_doc(doc)

        if isinstance(input, list):
            if tokenized:
                doc = self._list_to_document(input)
            else:
                if self._segmenter is None:
                    raise TypeError(
                        "Raw text input requires a segmenter. Initialize "
                        "COMBO with a language (e.g. COMBO('Polish')). "
                        "Or pass tokenized=True for pre-tokenized input."
                    )
                sentences = []
                for sent_text in input:
                    sent_doc = self._segmenter(
                        Document(sentences=[], text=sent_text), **seg_kwargs
                    )
                    for sent in sent_doc.sentences:
                        sent.sent_id = f"s{len(sentences) + 1}"
                        sentences.append(sent)
                doc = Document(sentences=sentences, text=" ".join(input))
            return self._predict_doc(doc)

        raise TypeError(f"Expected str, list, or Document, got {type(input).__name__}")

    def _predict_doc(self, doc: Document) -> Document:
        """Run the predictor on a segmented/tokenized Document."""
        if not doc.sentences:
            return doc

        # Separate empty sentences from those that need parsing
        non_empty_indices = []
        non_empty_sent_ids = []
        non_empty_tokens = []
        parsed_sentences: list[Optional[Sentence]] = [None] * len(doc.sentences)

        for i, sent in enumerate(doc.sentences):
            word_tokens = sent.tokens
            if not word_tokens:
                parsed_sentences[i] = sent
            else:
                non_empty_indices.append(i)
                non_empty_sent_ids.append(sent.sent_id)
                non_empty_tokens.append(word_tokens)

        if non_empty_tokens:
            predictions = self._predictor.predict_batch(
                [[t.form for t in tokens] for tokens in non_empty_tokens],
                sent_ids=non_empty_sent_ids,
                batch_size=self._batch_size,
                input_tokens_list=non_empty_tokens,
            )
            for idx, parsed in zip(non_empty_indices, predictions):
                parsed.text = doc.sentences[idx].text
                parsed_sentences[idx] = parsed

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

    @staticmethod
    def _list_to_document(input: "list[str] | list[list[str]]") -> Document:
        """Convert pre-tokenized input to a Document."""
        if not input:
            return Document(sentences=[], text="")

        # list[str] → single sentence; list[list[str]] → multiple sentences
        if isinstance(input[0], str):
            sentences_raw = [input]
        else:
            sentences_raw = input

        sentences = []
        for i, words in enumerate(sentences_raw):
            tokens = [Token(id=str(j + 1), form=w) for j, w in enumerate(words)]
            sentences.append(
                Sentence(
                    tokens=tokens,
                    sent_id=f"s{i + 1}",
                    text=" ".join(words),
                )
            )
        return Document(sentences=sentences, text="")

    def parse(self, words: list[str], sent_id: str = "s1") -> Sentence:
        """Parse a single pre-tokenized sentence.

        Args:
            words: List of word forms.
            sent_id: Sentence identifier.

        Returns:
            Sentence with all predictions.
        """
        return self._predictor.predict(words, sent_id=sent_id)

    def parse_batch(
        self,
        sentences: list[list[str]],
        sent_ids: Optional[list[str]] = None,
    ) -> list[Sentence]:
        """Parse multiple pre-tokenized sentences.

        Args:
            sentences: List of sentences, each a list of word forms.
            sent_ids: Optional sentence identifiers.

        Returns:
            List of parsed Sentences.
        """
        return self._predictor.predict_batch(
            sentences, sent_ids=sent_ids, batch_size=self._batch_size,
        )

    def __repr__(self) -> str:
        return f"COMBO(model={self._predictor.model.__class__.__name__}, batch_size={self._batch_size})"
