"""Tests for tracker-side layer-aware SafeTensors file selection.

The tracker advertises its local snapshot as a downloadable "model source" for
whatever layer range a node needs. If a required weight file for that range is
missing from the tracker's own disk, the tracker must refuse to advertise
itself as a source for that range — not quietly report the subset of files it
does have as if that were the complete, correct set. A partial-but-"complete"
answer makes the requesting node believe the download is already satisfied,
so it never fetches the real weights and only discovers the gap much later
at model-load time.
"""

import json

from meshnet_tracker.model_files import select_safetensors_files_for_layers


def _write_index(tmp_path, *, config=None):
    (tmp_path / "config.json").write_text(
        json.dumps(config or {"num_hidden_layers": 5}),
        encoding="utf-8",
    )
    (tmp_path / "model.safetensors.index.json").write_text(
        json.dumps({
            "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.4.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
                "model.norm.weight": "model-00004-of-00004.safetensors",
                "lm_head.weight": "model-00004-of-00004.safetensors",
            },
        }),
        encoding="utf-8",
    )


def _touch(path, size=1):
    path.write_bytes(b"0" * size)


def test_selects_files_when_snapshot_is_complete(tmp_path):
    "Selects files when snapshot is complete\n\nTags: general"
    _write_index(tmp_path)
    _touch(tmp_path / "model-00001-of-00004.safetensors")

    files = select_safetensors_files_for_layers(tmp_path, 0, 0)

    assert files == ["config.json", "model-00001-of-00004.safetensors", "model.safetensors.index.json"]


def test_returns_empty_when_a_required_weight_file_is_missing_on_disk(tmp_path):
    "Returns empty when a required weight file is missing on disk\n\nTags: general"
    _write_index(tmp_path)
    # model-00001-of-00004.safetensors is required for layer 0 but was never
    # downloaded onto this tracker host — the snapshot is incomplete.

    files = select_safetensors_files_for_layers(tmp_path, 0, 0)

    assert files == []


def test_tail_range_returns_empty_when_only_head_shard_present(tmp_path):
    "Tail range returns empty when only head shard present\n\nTags: general"
    _write_index(tmp_path)
    _touch(tmp_path / "model-00001-of-00004.safetensors")
    # model-00004-of-00004.safetensors (norm/lm_head) is required for the tail
    # range but missing — must not be silently dropped from the result.

    files = select_safetensors_files_for_layers(tmp_path, 4, 4)

    assert files == []