from __future__ import annotations

import json
from pathlib import Path

import pytest
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace

from deepinfer.hugging_face.tokenizer import get_hf_tokenizer

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]


@pytest.fixture
def checkpoint(tmp_path: Path) -> Path:
    tokenizer = Tokenizer(WordLevel({"[UNK]": 0, "hello": 1, "world": 2}, unk_token="[UNK]"))
    tokenizer.pre_tokenizer = Whitespace()
    tokenizer.save(str(tmp_path / "tokenizer.json"))
    (tmp_path / "config.json").write_text(json.dumps({"tokenizer_compatibility": "native"}))
    (tmp_path / "tokenizer_config.json").write_text(
        json.dumps(
            {
                "tokenizer_class": "CheckpointTokenizer",
                "unk_token": "[UNK]",
                "auto_map": {"AutoTokenizer": ["custom_tokenizer.CheckpointTokenizer", None]},
            }
        )
    )
    (tmp_path / "custom_tokenizer.py").write_text(
        "from pathlib import Path\n"
        "from transformers import PreTrainedTokenizerFast\n"
        "Path(__file__).with_suffix('.loaded').write_text('imported')\n"
        "class CheckpointTokenizer(PreTrainedTokenizerFast):\n"
        "    pass\n"
    )
    return tmp_path


@pytest.mark.parametrize("compatibility", ["native", "release"])
@pytest.mark.parametrize("explicit_false", [False, True], ids=["default", "false"])
def test_checkpoint_code_requires_trust(checkpoint: Path, compatibility: str, explicit_false: bool) -> None:
    if compatibility == "release":
        (checkpoint / "config.json").unlink()  # Exercise the default release policy.
    kwargs = {"trust_remote_code": False} if explicit_false else {}

    with pytest.raises(ValueError, match="trust_remote_code=True"):
        get_hf_tokenizer(str(checkpoint), **kwargs)

    assert not (checkpoint / "custom_tokenizer.loaded").exists()


def test_trusted_checkpoint_tokenizer_loads(checkpoint: Path) -> None:
    tokenizer = get_hf_tokenizer(str(checkpoint), trust_remote_code=True)

    assert (checkpoint / "custom_tokenizer.loaded").read_text() == "imported"
    assert type(tokenizer).__name__ == "CheckpointTokenizer"
    assert tokenizer.encode("hello world", add_special_tokens=False) == [1, 2]


@pytest.mark.parametrize("explicit_false", [False, True], ids=["default", "false"])
def test_builtin_tokenizer_does_not_require_trust(checkpoint: Path, explicit_false: bool) -> None:
    config_path = checkpoint / "tokenizer_config.json"
    config = json.loads(config_path.read_text())
    config.pop("auto_map")
    config["tokenizer_class"] = "PreTrainedTokenizerFast"
    config_path.write_text(json.dumps(config))
    kwargs = {"trust_remote_code": False} if explicit_false else {}

    tokenizer = get_hf_tokenizer(str(checkpoint), **kwargs)

    assert tokenizer.encode("hello world", add_special_tokens=False) == [1, 2]
    assert not (checkpoint / "custom_tokenizer.loaded").exists()