from collections.abc import Iterator

import pytest
from tokenizers import Tokenizer

from skeletoken.base import TokenizerModel
from skeletoken.merges import Merges
from skeletoken.models import BPE, ModelType, Unigram, WordPiece
from skeletoken.trainer import train_tokenizer
from skeletoken.vocabulary import UnigramVocabulary, Vocabulary


def _dummy_data() -> Iterator[str]:
    """
    Small iterator of strings used as fake training data.

    Returns
    -------
    Iterator[str]

    """
    for s in ["hello world", "another sentence"]:
        yield s


@pytest.mark.parametrize("model_type", [ModelType.BPE, ModelType.UNIGRAM, ModelType.WORDPIECE])
def test_train_calls_tokenizer_train_from_iterator(model_type: ModelType, monkeypatch) -> None:
    """
    Ensure train_tokenizer calls Tokenizer.train_from_iterator with a trainer for the given model type.

    We mock Tokenizer.train_from_iterator to be a no-op and assert it was called with a trainer instance.
    """
    # Construct a tiny vocabulary usable for all models
    vocab = {"[PAD]": 0, "[UNK]": 1, "a": 2, " ": 3}

    model: BPE | Unigram | WordPiece
    if model_type == ModelType.BPE:
        model = BPE(
            merges=Merges([]),
            vocab=Vocabulary(vocab),
            dropout=0.0,
            unk_token="[UNK]",
            continuing_subword_prefix="",
            end_of_word_suffix="",
            fuse_unk=False,
            byte_fallback=False,
            ignore_merges=False,
        )
    elif model_type == ModelType.UNIGRAM:
        u_root = [(t, -1.0) for t in ["[PAD]", "[UNK]", "a", " "]]
        model = Unigram(vocab=UnigramVocabulary(u_root), unk_id=1, byte_fallback=False)
    else:
        model = WordPiece(vocab=Vocabulary(vocab), unk_token="[UNK]", continuing_subword_prefix="")

    tokenizer_model = TokenizerModel(model=model)
    tokenizer = tokenizer_model.to_tokenizer()
    called: dict = {}

    def fake_train_from_iterator(data: Iterator[str], trainer) -> None:
        called["data"] = list(data)
        called["trainer_type"] = type(trainer).__name__

    monkeypatch.setattr(tokenizer, "train_from_iterator", fake_train_from_iterator)

    # Monkeypatch TokenizerModel.to_tokenizer to return our tokenizer
    monkeypatch.setattr(TokenizerModel, "to_tokenizer", lambda self: tokenizer)

    out = train_tokenizer(tokenizer_model, _dummy_data(), vocab_size=100)

    assert "trainer_type" in called
    assert isinstance(out, TokenizerModel)


def test_train_unknown_model_raises() -> None:
    """If model type is unknown, train_tokenizer should raise ValueError."""

    # Create a fake object with the minimal interface used by train_tokenizer
    class Fake:
        def __init__(self) -> None:
            class M:
                type = "Unknown"

            self.model = M()

            class AT:
                def get_special_tokens(self) -> list:
                    return []

            self.added_tokens = AT()

        def to_tokenizer(self) -> Tokenizer:
            return Tokenizer.from_file("tests/data/wordpiece/tokenizer.json")

    fake = Fake()

    with pytest.raises(ValueError):
        train_tokenizer(fake, _dummy_data(), vocab_size=10)  # type: ignore
