"""
Tests for generate_dataset_hash_from_config.

Regression test for https://github.com/axolotl-ai-cloud/axolotl/issues/3303:
changing output_dir should not bust the dataset cache when added_tokens_overrides
is set.
"""

from axolotl.utils.data.shared import generate_dataset_hash_from_config
from axolotl.utils.dict import DictDefault


def _base_cfg(**kwargs):
    return DictDefault(
        {
            "sequence_len": 2048,
            "sample_packing": False,
            "eval_sample_packing": False,
            "group_by_length": False,
            "kd_temperature": None,
            "dataset_exact_deduplication": False,
            "tokenizer_config": "NousResearch/Llama-3.2-1B",
            **kwargs,
        }
    )


def _datasets():
    return [
        DictDefault(
            {
                "path": "mhenrichsen/alpaca_2k_test",
                "type": "alpaca",
                "shards": None,
                "conversation": None,
                "split": "train",
                "temperature": None,
            }
        )
    ]


class TestGenerateDatasetHashFromConfig:
    def test_same_config_same_hash(self):
        """Identical configs produce identical hashes."""
        cfg = _base_cfg()
        h1 = generate_dataset_hash_from_config(
            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
        )
        h2 = generate_dataset_hash_from_config(
            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
        )
        assert h1 == h2

    def test_different_tokenizer_different_hash(self):
        """A different tokenizer path produces a different hash."""
        cfg = _base_cfg()
        h1 = generate_dataset_hash_from_config(
            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
        )
        h2 = generate_dataset_hash_from_config(
            cfg, _datasets(), "HuggingFaceTB/SmolLM2-135M"
        )
        assert h1 != h2

    def test_different_sequence_len_different_hash(self):
        cfg_a = _base_cfg(sequence_len=2048)
        cfg_b = _base_cfg(sequence_len=4096)
        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")
        assert h1 != h2

    def test_different_special_tokens_different_hash(self):
        cfg_a = _base_cfg(special_tokens={"pad_token": "<|endoftext|>"})
        cfg_b = _base_cfg(
            special_tokens={
                "pad_token": "<|endoftext|>",
                "bos_token": "<|custom_im_start|>",
                "eos_token": "<|custom_im_end|>",
            }
        )

        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")

        assert h1 != h2

    def test_different_chat_template_different_hash(self):
        cfg_a = _base_cfg(chat_template="chatml")
        cfg_b = _base_cfg(chat_template="jinja", chat_template_jinja="{{ messages }}")

        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")

        assert h1 != h2

    def test_dataset_chat_template_fields_affect_hash(self):
        datasets_a = _datasets()
        datasets_b = _datasets()
        datasets_b[0].chat_template = "jinja"
        datasets_b[0].chat_template_jinja = "{{ messages }}"
        datasets_b[0].field_messages = "conversations"
        datasets_b[0].message_property_mappings = {
            "role": "from",
            "content": "value",
        }

        cfg = _base_cfg()
        h1 = generate_dataset_hash_from_config(cfg, datasets_a, "tok")
        h2 = generate_dataset_hash_from_config(cfg, datasets_b, "tok")

        assert h1 != h2

    def test_llama_fix_token_configs_do_not_share_cache_hash(self):
        datasets_custom_tokens = [
            DictDefault(
                {
                    "path": "mlabonne/FineTome-100k",
                    "type": "chat_template",
                    "split": "train[:10%]",
                    "chat_template": "jinja",
                    "chat_template_jinja": "{{ messages }}",
                    "field_messages": "conversations",
                    "message_property_mappings": {
                        "role": "from",
                        "content": "value",
                    },
                }
            )
        ]
        datasets_chatml = [
            DictDefault(
                {
                    "path": "mlabonne/FineTome-100k",
                    "type": "chat_template",
                    "split": "train[:10%]",
                    "field_messages": "conversations",
                    "message_property_mappings": {
                        "role": "from",
                        "content": "value",
                    },
                }
            )
        ]
        cfg_custom_tokens = _base_cfg(
            sequence_len=512,
            sample_packing=True,
            eval_sample_packing=True,
            special_tokens={
                "pad_token": "<|endoftext|>",
                "bos_token": "<|custom_im_start|>",
                "eos_token": "<|custom_im_end|>",
            },
        )
        cfg_chatml = _base_cfg(
            sequence_len=512,
            sample_packing=True,
            eval_sample_packing=True,
            special_tokens={"pad_token": "<|endoftext|>"},
            chat_template="chatml",
        )

        h1 = generate_dataset_hash_from_config(
            cfg_custom_tokens, datasets_custom_tokens, "HuggingFaceTB/SmolLM2-135M"
        )
        h2 = generate_dataset_hash_from_config(
            cfg_chatml, datasets_chatml, "HuggingFaceTB/SmolLM2-135M"
        )

        assert h1 != h2

    # --- Regression: added_tokens_overrides + output_dir ---

    def test_added_tokens_overrides_hash_stable_across_output_dir(self):
        """Hash must not change when only output_dir changes (issue #3303).

        When added_tokens_overrides is set the tokenizer is saved into output_dir,
        making tokenizer.name_or_path an absolute path that includes output_dir.
        The hash should be derived from the canonical tokenizer config + overrides,
        not from the output-dir-dependent path.
        """
        cfg_run1 = _base_cfg(
            output_dir="/tmp/run_1",
            added_tokens_overrides={32000: "<PAD>", 32001: "<MASK>"},
        )
        cfg_run2 = _base_cfg(
            output_dir="/tmp/run_2_different_name",
            added_tokens_overrides={32000: "<PAD>", 32001: "<MASK>"},
        )

        # Simulate what happens in practice: tokenizer.name_or_path becomes the
        # output_dir-based path after modify_tokenizer_files() saves the tokenizer.
        tokenizer_name_run1 = "/tmp/run_1/modified_tokenizer"
        tokenizer_name_run2 = "/tmp/run_2_different_name/modified_tokenizer"

        h1 = generate_dataset_hash_from_config(
            cfg_run1, _datasets(), tokenizer_name_run1
        )
        h2 = generate_dataset_hash_from_config(
            cfg_run2, _datasets(), tokenizer_name_run2
        )

        assert h1 == h2, (
            "Dataset cache hash must not change when only output_dir changes "
            "while added_tokens_overrides stays the same (issue #3303)."
        )

    def test_added_tokens_overrides_different_overrides_different_hash(self):
        """Different added_tokens_overrides produce different hashes."""
        cfg_a = _base_cfg(
            output_dir="/tmp/run_a",
            added_tokens_overrides={32000: "<PAD>"},
        )
        cfg_b = _base_cfg(
            output_dir="/tmp/run_a",  # same output_dir
            added_tokens_overrides={32000: "<OTHER>"},
        )
        tokenizer_path = "/tmp/run_a/modified_tokenizer"

        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), tokenizer_path)
        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), tokenizer_path)

        assert h1 != h2

    def test_no_added_tokens_overrides_uses_tokenizer_name_as_before(self):
        """Without added_tokens_overrides the old behaviour is preserved."""
        cfg = _base_cfg()  # no added_tokens_overrides
        tokenizer_name = "NousResearch/Llama-3.2-1B"

        h1 = generate_dataset_hash_from_config(cfg, _datasets(), tokenizer_name)
        # Changing tokenizer_name still changes the hash
        h2 = generate_dataset_hash_from_config(cfg, _datasets(), "some/other-model")

        assert h1 != h2