# Partly Copyright (C) 2025 Arcee AI
# Partly Copyright (C) 2025-2026 Allura-org

from typing import Dict, List, Optional

import pytest
import torch
from common import make_picollama, make_tokenizer, run_and_check_merge
from transformers import LlamaConfig, LlamaTokenizerFast

from mergekitty.config import InputModelDefinition, MergeConfiguration
from mergekitty.io import LazyTensorLoader
from mergekitty.tokenizer import TokenizerConfig


@pytest.fixture(scope="session")
def model_base(tmp_path_factory):
    model_path = make_picollama(tmp_path_factory.mktemp("model_base"), vocab_size=64)
    make_tokenizer(vocab_size=64, added_tokens=[]).save_pretrained(model_path)
    return model_path


@pytest.fixture(scope="session")
def model_chatml(tmp_path_factory):
    model_path = make_picollama(tmp_path_factory.mktemp("model_chatml"), vocab_size=66)
    make_tokenizer(
        vocab_size=64, added_tokens=["<|im_start|>", "<|im_end|>"]
    ).save_pretrained(model_path)
    return model_path


@pytest.fixture(scope="session")
def model_padded(tmp_path_factory):
    model_path = make_picollama(tmp_path_factory.mktemp("model_padded"), vocab_size=64)
    make_tokenizer(
        vocab_size=64,
        added_tokens=["<UNUSED_0>", "<UNUSED_1>", "<UNUSED_2>", "<UNUSED_3>"],
    ).save_pretrained(model_path)
    return model_path


def check_tokenizer(
    expected_size: int,
    expected_added_ct: Optional[int] = None,
    must_contain: Optional[List[str]] = None,
    must_not_contain: Optional[List[str]] = None,
):
    def _cb(model_path: str):
        tok: LlamaTokenizerFast = LlamaTokenizerFast.from_pretrained(model_path)
        vocab = tok.get_vocab()
        print(vocab)
        assert len(vocab) == expected_size

        if expected_added_ct is not None:
            assert len(tok.added_tokens_decoder) == expected_added_ct

        if must_contain:
            for tok in must_contain:
                assert tok in vocab

        if must_not_contain:
            for tok in must_not_contain:
                assert tok not in vocab

    return _cb


class ModelEmbeddings:
    embed_tokens: torch.Tensor
    vocab: Dict[str, int]

    def __init__(self, model_path: str):
        tokenizer = LlamaTokenizerFast.from_pretrained(model_path)
        loader = LazyTensorLoader.from_disk(model_path)
        self.embed_tokens = loader.get_tensor("model.embed_tokens.weight")
        self.vocab = tokenizer.get_vocab()

    def token_embedding(self, token: str) -> Optional[torch.Tensor]:
        idx = self.vocab.get(token)
        if idx is None:
            return None
        return self.embed_tokens[idx, :]


class TestTokenizerMerges:
    def test_legacy_mode(self, model_base: str, model_padded: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_padded, model_chatml], base_model=model_base
        )
        # when no tokenizer_source is set, expect output tokenizer to be from base_model
        run_and_check_merge(
            config, validate=check_tokenizer(expected_size=64, expected_added_ct=3)
        )

    def test_source_base(self, model_base: str, model_padded: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_padded, model_chatml],
            base_model=model_base,
            tokenizer_source="base",
        )
        # expect the same output but it's a different code path
        run_and_check_merge(
            config, validate=check_tokenizer(expected_size=64, expected_added_ct=3)
        )

    def test_source_union(self, model_base: str, model_padded: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_padded, model_chatml],
            base_model=model_base,
            tokenizer_source="union",
        )

        def _check_embed(model_path: str):
            # output should have all tokens used by any model
            # but not include any unused tokens
            check_tokenizer(
                expected_size=66,
                expected_added_ct=5,
                must_contain=["<|im_start|>", "<|im_end|>"],
                must_not_contain=[f"<UNUSED_{idx}>" for idx in range(4)],
            )(model_path)
            emb_out = ModelEmbeddings(model_path)
            emb_chatml = ModelEmbeddings(model_chatml)

            assert torch.allclose(
                emb_out.token_embedding("<|im_start|>"),
                emb_chatml.token_embedding("<|im_start|>"),
            ), "Token <|im_start|> should be from model_chatml"
            assert torch.allclose(
                emb_out.token_embedding("<|im_end|>"),
                emb_chatml.token_embedding("<|im_end|>"),
                atol=1e-3,
                rtol=1e-4,
            ), "Token <|im_end|> should be from model_chatml"

        run_and_check_merge(
            config,
            validate=_check_embed,
        )

    def test_source_model(self, model_base: str, model_padded: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_padded, model_chatml],
            base_model=model_base,
            tokenizer_source=model_chatml,
        )
        # tokenizer should match model_chatml
        run_and_check_merge(
            config,
            validate=check_tokenizer(
                expected_size=66, must_contain=["<|im_start|>", "<|im_end|>"]
            ),
        )

    def test_slerp_union(self, model_base: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_chatml],
            base_model=model_base,
            tokenizer_source="union",
            merge_method="slerp",
            t=0.5,
        )

        run_and_check_merge(
            config,
            validate=check_tokenizer(
                expected_size=66,
                must_contain=["<|im_start|>", "<|im_end|>"],
            ),
        )

    def test_force_token(self, model_base: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_chatml],
            base_model=model_base,
            merge_method="linear",
            tokenizer_config=TokenizerConfig(
                source="union",
                tokens={
                    "_tok_10": {"source": model_chatml, "force": True},
                    "_tok_11": {"source": model_base, "force": True},
                },
            ),
        )

        def _check_embed(model_path: str):
            check_tokenizer(
                expected_size=66, must_contain=["<|im_start|>", "<|im_end|>"]
            )(model_path)
            emb_out = ModelEmbeddings(model_path)
            emb_base = ModelEmbeddings(model_base)
            emb_chatml = ModelEmbeddings(model_chatml)

            assert torch.allclose(
                emb_out.token_embedding("_tok_10"),
                emb_chatml.token_embedding("_tok_10"),
            ), "Token _tok_10 should be from model_chatml"
            assert torch.allclose(
                emb_out.token_embedding("_tok_11"),
                emb_base.token_embedding("_tok_11"),
            ), "Token _tok_11 should be from model_base"

        run_and_check_merge(config, validate=_check_embed)

    def test_model_token_id(self, model_base: str, model_chatml: str):
        config = self.make_config(
            [model_base, model_chatml],
            base_model=model_base,
            merge_method="linear",
            tokenizer_config=TokenizerConfig(
                source="base",
                tokens={
                    "_tok_20": {
                        "source": {
                            "kind": "model_token",
                            "model": model_chatml,
                            "token_id": 64,
                        },
                        "force": True,
                    },
                    "_tok_21": {
                        "source": {
                            "kind": "model_token",
                            "model": model_base,
                            "token": "<s>",
                        },
                        "force": True,
                    },
                },
            ),
        )

        def _check_embed(model_path: str):
            check_tokenizer(expected_size=64, must_contain=["_tok_10"])(model_path)
            emb_out = ModelEmbeddings(model_path)
            emb_base = ModelEmbeddings(model_base)
            emb_chatml = ModelEmbeddings(model_chatml)

            assert torch.allclose(
                emb_out.token_embedding("_tok_20"), emb_chatml.embed_tokens[64, :]
            ), "Token _tok_20 should be == model_chatml token 64"
            assert torch.allclose(
                emb_out.token_embedding("_tok_21"), emb_base.token_embedding("<s>")
            ), "Token _tok_21 should be == model_base <s>"

        run_and_check_merge(config, validate=_check_embed)

    def test_pad_to_multiple_of(self, model_chatml: str):
        config = self.make_config(
            [model_chatml],
            base_model=model_chatml,
            merge_method="linear",
            tokenizer_config=TokenizerConfig(
                source="base",
                pad_to_multiple_of=16,
            ),
        )
        real_vocab_size = 64 + 2
        padded_size = (real_vocab_size // 16 + 1) * 16

        def _check_result(model_path: str):
            cfg = LlamaConfig.from_pretrained(model_path)
            assert cfg.vocab_size == padded_size, (
                f"Expected vocab size {padded_size}, got {cfg.vocab_size}"
            )
            check_tokenizer(
                expected_size=real_vocab_size,
                must_contain=["<|im_start|>", "<|im_end|>"],
            )(model_path)

            emb_out = ModelEmbeddings(model_path)
            assert emb_out.embed_tokens.shape[0] == padded_size, (
                "Embedding size mismatch"
            )

        run_and_check_merge(config, validate=_check_result)

    def make_config(
        self,
        models: List[str],
        base_model: Optional[str] = None,
        merge_method: str = "linear",
        tokenizer_source: Optional[str] = None,
        t: Optional[float] = None,
        tokenizer_config: Optional[TokenizerConfig] = None,
    ):
        parameters = {"t": t} if t is not None else {}
        config = MergeConfiguration(
            merge_method=merge_method,
            base_model=base_model,
            models=[
                InputModelDefinition(
                    model=m,
                    parameters={"weight": 1.0},
                )
                for m in models
            ],
            dtype="float32",
            tokenizer_source=tokenizer_source,
            parameters=parameters,
            tokenizer=tokenizer_config,
        )
        return config
