"""Tests for layer-aware SafeTensors snapshot file selection."""

import json

import pytest

from meshnet_node.safetensors_selection import select_safetensors_files_for_layers


def _write_snapshot(tmp_path, *, config=None):
    (tmp_path / "config.json").write_text(
        json.dumps(config or {"num_hidden_layers": 5}),
        encoding="utf-8",
    )
    (tmp_path / "tokenizer.json").write_text("{}", encoding="utf-8")
    (tmp_path / "tokenizer_config.json").write_text("{}", encoding="utf-8")
    (tmp_path / "README.md").write_text("not part of runtime snapshot", encoding="utf-8")
    (tmp_path / "model.safetensors.index.json").write_text(
        json.dumps({
            "metadata": {"total_size": 123},
            "weight_map": {
                "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
                "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
                "model.layers.1.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
                "model.layers.2.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
                "model.layers.3.mlp.down_proj.weight": "nested/model-00003-of-00004.safetensors",
                "model.layers.4.self_attn.q_proj.weight": "nested/model-00003-of-00004.safetensors",
                "model.norm.weight": "model-00004-of-00004.safetensors",
                "lm_head.weight": "model-00004-of-00004.safetensors",
            },
        }),
        encoding="utf-8",
    )


def test_selects_only_weight_shards_for_middle_layer_range(tmp_path):
    _write_snapshot(tmp_path)

    files = select_safetensors_files_for_layers(tmp_path, 2, 3)

    assert files == [
        "config.json",
        "model-00002-of-00004.safetensors",
        "model.safetensors.index.json",
        "nested/model-00003-of-00004.safetensors",
        "tokenizer.json",
        "tokenizer_config.json",
    ]


def test_head_range_includes_embeddings(tmp_path):
    _write_snapshot(tmp_path)

    files = select_safetensors_files_for_layers(tmp_path, 0, 0)

    assert "model-00001-of-00004.safetensors" in files
    assert "model-00004-of-00004.safetensors" not in files


def test_tail_range_includes_norm_and_lm_head_from_inferred_layer_count(tmp_path):
    _write_snapshot(tmp_path, config={"text_config": {"num_hidden_layers": 5}})

    files = select_safetensors_files_for_layers(tmp_path, 4, 4)

    assert "nested/model-00003-of-00004.safetensors" in files
    assert "model-00004-of-00004.safetensors" in files
    assert "model-00001-of-00004.safetensors" not in files


def test_tail_files_are_not_selected_without_total_layer_count(tmp_path):
    _write_snapshot(tmp_path, config={"architectures": ["UnknownForTest"]})

    files = select_safetensors_files_for_layers(tmp_path, 4, 4)

    assert "nested/model-00003-of-00004.safetensors" in files
    assert "model-00004-of-00004.safetensors" not in files


def test_rejects_unsafe_weight_map_paths(tmp_path):
    (tmp_path / "model.safetensors.index.json").write_text(
        json.dumps({"weight_map": {"model.layers.0.weight": "../escape.safetensors"}}),
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="unsafe relative file"):
        select_safetensors_files_for_layers(tmp_path, 0, 0)