"""Dense-Llama GGUF ownership selection and introspection tests."""

from __future__ import annotations

import pytest

from meshnet_node.gguf_ownership import (
    DenseLlamaShardOwnership,
    authoritative_dense_llama_ownership,
    infer_dense_llama_ownership,
    select_dense_llama_tensor_names,
)


def test_dense_llama_selection_only_picks_block_range_and_endpoints():
    "Dense-Llama selection keeps only the owned blocks plus the correct endpoints.\n\nTags: node, GGUF"
    tensor_inventory = {
        "token_embd.weight": 10_000,
        "blk.0.attn_q.weight": 1_000,
        "blk.0.ffn_down.weight": 1_000,
        "blk.1.attn_q.weight": 2_000,
        "blk.1.ffn_down.weight": 2_000,
        "blk.2.attn_q.weight": 3_000,
        "blk.2.ffn_down.weight": 3_000,
        "output_norm.weight": 256,
        "output.weight": 10_000,
        "rope.freqs": 128,
    }

    selected = select_dense_llama_tensor_names(
        tensor_inventory,
        1,
        2,
        total_layers=3,
    )

    assert selected == {
        "blk.1.attn_q.weight",
        "blk.1.ffn_down.weight",
        "blk.2.attn_q.weight",
        "blk.2.ffn_down.weight",
        "output_norm.weight",
        "output.weight",
    }

    selected_bytes = sum(tensor_inventory[name] for name in selected)
    full_bytes = sum(tensor_inventory.values())
    assert selected_bytes == 20_256
    assert selected_bytes < full_bytes


def test_dense_llama_loaded_range_is_authoritative_from_tensor_inventory():
    "The backend's loaded tensor inventory is the source of truth for range and ownership.\n\nTags: node, GGUF"

    class Backend:
        loaded_tensor_names = (
            "token_embd.weight",
            "blk.4.attn_q.weight",
            "blk.5.ffn_down.weight",
            "output_norm.weight",
            "output.weight",
        )

    ownership = authoritative_dense_llama_ownership(Backend(), selection=None)

    assert isinstance(ownership, DenseLlamaShardOwnership)
    assert ownership.range == (4, 5)
    assert ownership.owns_embedding is True
    assert ownership.owns_final_head is True


def test_derivative_slice_requires_source_and_slice_hashes():
    "Temporary derivative GGUF slices must carry hashes and cannot claim final semantics.\n\nTags: node, GGUF"
    with pytest.raises(ValueError, match="source and slice hashes"):
        infer_dense_llama_ownership(
            ["blk.1.attn_q.weight"],
            derivative_slice=True,
            final_artifact_semantics=False,
        )

    with pytest.raises(ValueError, match="final artifacts"):
        infer_dense_llama_ownership(
            ["blk.1.attn_q.weight"],
            source_artifact_hash="sha256:source",
            slice_artifact_hash="sha256:slice",
            derivative_slice=True,
            final_artifact_semantics=True,
        )