"""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", "."])
    result = nlp([["Ala", "ma", "kota", "."], ["Pies", "je", "."]])

    # 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 .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_or_predictor, device: Optional[str] = None):
        """Initialize COMBO with a language or a Predictor instance.

        Args:
            language_or_predictor: Language enum member (e.g. Language.POLISH),
                language name string (e.g. "Polish"),
                or a Predictor instance directly.
            device: Device string ("cuda", "cpu", "mps", "auto"). Only used
                when passing a language. Auto-detected if None.

        Example::

            from combo import COMBO, Language
            nlp = COMBO(Language.POLISH)
            nlp = COMBO("Polish")
            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

        if isinstance(language_or_predictor, (str, Language)):
            lang = resolve_language(language_or_predictor)
            from .hub import resolve_model, resolve_lambo_model
            model_id = resolve_model(lang.value)
            instance = COMBO.from_pretrained(model_id, device=device)
            self._predictor = instance._predictor

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

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        device: Optional[str] = None,
    ) -> "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,
        )

        return cls(predictor)

    def __call__(self, input: "str | list[str] | list[list[str]] | Document") -> Document:
        """Parse input text, pre-tokenized sentence(s), or Document.

        Args:
            input:
                - ``str``: Raw text — requires a segmenter (use ``COMBO(language)``).
                - ``list[str]``: Single pre-tokenized sentence.
                - ``list[list[str]]``: Multiple pre-tokenized sentences.
                - ``Document``: From a previous pipeline component.

        Returns:
            Document with parsed sentences.
        """
        if isinstance(input, str):
            if self._segmenter is None:
                raise TypeError(
                    "Raw text input requires a segmenter. Either initialize "
                    "COMBO with a language (e.g. COMBO('Polish')) or pass "
                    "pre-tokenized input as a list of strings."
                )
            doc = self._segmenter(Document(sentences=[], text=input))
        elif isinstance(input, list):
            doc = self._list_to_document(input)
        elif isinstance(input, Document):
            doc = input
        else:
            raise TypeError(f"Expected str, list, or Document, got {type(input).__name__}")

        if not doc.sentences:
            return doc

        # Parse each sentence
        parsed_sentences = []
        for sent in doc.sentences:
            words = [t.form for t in sent.words]
            if not words:
                parsed_sentences.append(sent)
                continue
            parsed = self._predictor.predict(words, sent_id=sent.sent_id)
            parsed.text = sent.text
            parsed_sentences.append(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)

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