"""DGR-034: end-to-end owned-range loads through the native report tool.

Gated on the built ``meshnet-range-report`` binary (the deterministic
CPU-only native lane builds it from the pinned, patched llama.cpp tree); in
an environment without that build these tests skip rather than fake a pass.
When the binary is present they run real loads of a tiny synthetic
dense-Llama GGUF — no model download, no GPU — and prove the loader
registers exactly the owned tensors, reports ownership derived from the
loaded state, and rejects invalid/out-of-model ranges and missing required
tensors. The JSON is consumed through ``meshnet_node.range_report`` so the
strict project-owned contract is exercised on real tool output.
"""

from __future__ import annotations

import json
import os
import struct
import subprocess
import sys
from pathlib import Path

import pytest

from meshnet_node.range_report import RangeReportError, parse_owned_range_report

REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_BINARY = REPO_ROOT / "build" / "llama.cpp" / "build" / "bin" / "meshnet-range-report"

BINARY = Path(os.environ.get("MESHNET_RANGE_REPORT_BIN", DEFAULT_BINARY))

requires_range_report_tool = pytest.mark.skipif(
    not BINARY.is_file(),
    reason=(
        "meshnet-range-report is not built; run the deterministic native lane "
        "(scripts/llama_cpp_dependency.py build) to enable these tests"
    ),
)

# --- Minimal GGUF v3 writer, mirroring the model-free native fixture --------

K_LAYERS = 4
K_EMBD = 8
K_FFN = 16
K_VOCAB = 16
ALIGNMENT = 32

_GGUF_UINT32 = 4
_GGUF_FLOAT32 = 6
_GGUF_STRING = 8
_GGML_TYPE_F32 = 0


def _gguf_string(value: str) -> bytes:
    data = value.encode("utf-8")
    return struct.pack("<Q", len(data)) + data


def _metadata_entries() -> list[tuple[str, int, object]]:
    return [
        ("general.architecture", _GGUF_STRING, "llama"),
        ("general.alignment", _GGUF_UINT32, ALIGNMENT),
        ("llama.context_length", _GGUF_UINT32, 16),
        ("llama.embedding_length", _GGUF_UINT32, K_EMBD),
        ("llama.block_count", _GGUF_UINT32, K_LAYERS),
        ("llama.feed_forward_length", _GGUF_UINT32, K_FFN),
        ("llama.attention.head_count", _GGUF_UINT32, 2),
        ("llama.attention.head_count_kv", _GGUF_UINT32, 2),
        ("llama.rope.dimension_count", _GGUF_UINT32, 4),
        ("llama.attention.layer_norm_rms_epsilon", _GGUF_FLOAT32, 1.0e-5),
        ("tokenizer.ggml.model", _GGUF_STRING, "no_vocab"),
        ("llama.vocab_size", _GGUF_UINT32, K_VOCAB),
    ]


def _fixture_tensors() -> list[tuple[str, tuple[int, ...]]]:
    tensors: list[tuple[str, tuple[int, ...]]] = [
        ("token_embd.weight", (K_EMBD, K_VOCAB)),
        ("output_norm.weight", (K_EMBD,)),
        ("output.weight", (K_EMBD, K_VOCAB)),
    ]
    for layer in range(K_LAYERS):
        prefix = f"blk.{layer}."
        tensors += [
            (prefix + "attn_norm.weight", (K_EMBD,)),
            (prefix + "attn_q.weight", (K_EMBD, K_EMBD)),
            (prefix + "attn_k.weight", (K_EMBD, K_EMBD)),
            (prefix + "attn_v.weight", (K_EMBD, K_EMBD)),
            (prefix + "attn_output.weight", (K_EMBD, K_EMBD)),
            (prefix + "ffn_norm.weight", (K_EMBD,)),
            (prefix + "ffn_gate.weight", (K_EMBD, K_FFN)),
            (prefix + "ffn_down.weight", (K_FFN, K_EMBD)),
            (prefix + "ffn_up.weight", (K_EMBD, K_FFN)),
        ]
    return tensors


def write_dense_llama_gguf(path: Path, *, drop: frozenset[str] = frozenset()) -> Path:
    """Write a tiny dense-Llama GGUF; ``drop`` omits tensors (corruption cases)."""
    kvs = _metadata_entries()
    tensors = [(name, dims) for name, dims in _fixture_tensors() if name not in drop]

    blob = bytearray()
    blob += b"GGUF" + struct.pack("<IQQ", 3, len(tensors), len(kvs))
    for key, vtype, value in kvs:
        blob += _gguf_string(key)
        blob += struct.pack("<I", vtype)
        if vtype == _GGUF_STRING:
            blob += _gguf_string(value)  # type: ignore[arg-type]
        elif vtype == _GGUF_UINT32:
            blob += struct.pack("<I", value)  # type: ignore[arg-type]
        elif vtype == _GGUF_FLOAT32:
            blob += struct.pack("<f", value)  # type: ignore[arg-type]
        else:  # pragma: no cover - writer guard
            raise AssertionError(f"unhandled kv type {vtype}")

    offset = 0
    infos = bytearray()
    data = bytearray()
    for name, dims in tensors:
        infos += _gguf_string(name)
        infos += struct.pack("<I", len(dims))
        for dim in dims:
            infos += struct.pack("<Q", dim)
        infos += struct.pack("<IQ", _GGML_TYPE_F32, offset)
        size = 4
        for dim in dims:
            size *= dim
        assert size % ALIGNMENT == 0
        data += bytes(size)
        offset += size

    blob += infos
    blob += bytes(-len(blob) % ALIGNMENT)  # pad header to the data section
    blob += data
    path.write_bytes(bytes(blob))
    return path


# --- Tool driver -------------------------------------------------------------

LAYER_BYTES = 2624  # 9 registered F32 tensors per layer, see _fixture_tensors
EMBD_BYTES = 512
OUT_NORM_BYTES = 32
OUT_BYTES = 512


def run_tool(model: Path, start: int, end: int, *extra: str) -> tuple[int, dict]:
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = f"{BINARY.parent}:{env.get('LD_LIBRARY_PATH', '')}"
    completed = subprocess.run(
        [
            str(BINARY),
            "--model", str(model),
            "--start", str(start),
            "--end", str(end),
            *extra,
        ],
        capture_output=True,
        text=True,
        env=env,
        timeout=120,
    )
    try:
        doc = json.loads(completed.stdout)
    except json.JSONDecodeError as exc:  # pragma: no cover - diagnostic path
        raise AssertionError(
            f"tool did not print a JSON report (exit {completed.returncode}): "
            f"{completed.stdout!r} {completed.stderr!r}"
        ) from exc
    return completed.returncode, doc


@pytest.fixture(scope="module")
def dense_llama_gguf(tmp_path_factory: pytest.TempPathFactory) -> Path:
    return write_dense_llama_gguf(tmp_path_factory.mktemp("gguf") / "dense-llama.gguf")


@requires_range_report_tool
class TestOwnedRangeLoads:
    def test_middle_range_registers_exactly_its_layers(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 1, 3, "--no-extra-bufts")
        assert code == 0
        report = parse_owned_range_report(doc)
        assert (report.start_layer, report.end_layer) == (1, 3)
        assert report.registered_tensors == 18
        assert report.registered_bytes == 2 * LAYER_BYTES
        # The fixture layers are contiguous in the file, so the pure mmap span
        # is exactly the owned tensor bytes — scaled down from the artifact.
        assert report.mapped_bytes == 2 * LAYER_BYTES
        assert report.mapped_bytes < report.file_bytes

    def test_head_range_owns_embeddings(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 0, 1)
        assert code == 0
        report = parse_owned_range_report(doc)
        assert report.is_head and report.has_token_embeddings
        assert not report.has_output_head
        assert report.registered_tensors == 10
        assert report.registered_bytes == EMBD_BYTES + LAYER_BYTES

    def test_tail_range_owns_norm_and_output(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 3, 4)
        assert code == 0
        report = parse_owned_range_report(doc)
        assert report.is_tail and report.has_output_head
        assert not report.has_token_embeddings
        assert report.registered_tensors == 11
        assert report.registered_bytes == LAYER_BYTES + OUT_NORM_BYTES + OUT_BYTES

    def test_shards_partition_the_whole_model_bytes(self, dense_llama_gguf: Path) -> None:
        shards = [(0, 1), (1, 3), (3, 4)]
        registered = []
        for start, end in shards:
            code, doc = run_tool(dense_llama_gguf, start, end)
            assert code == 0
            registered.append(parse_owned_range_report(doc).registered_bytes)
        code, doc = run_tool(dense_llama_gguf, 0, 4)
        assert code == 0
        whole = parse_owned_range_report(doc)
        assert whole.registered_tensors == 3 + 9 * K_LAYERS
        assert sum(registered) == whole.registered_bytes

    def test_non_mmap_load_scales_resident_with_the_range(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 1, 3, "--no-mmap")
        assert code == 0
        report = parse_owned_range_report(doc)
        assert report.mapped_bytes == 0
        assert report.registered_bytes == 2 * LAYER_BYTES
        code, doc = run_tool(dense_llama_gguf, 0, 4, "--no-mmap")
        assert code == 0
        whole = parse_owned_range_report(doc)
        assert report.resident_bytes < whole.resident_bytes


@requires_range_report_tool
class TestRangeRejection:
    def test_out_of_model_range_is_refused(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 3, 5)
        assert code == 3 and doc["ok"] is False
        with pytest.raises(RangeReportError):
            parse_owned_range_report(doc)

    def test_empty_range_is_refused(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 2, 2)
        assert code == 3 and doc["ok"] is False

    def test_inverted_range_is_refused(self, dense_llama_gguf: Path) -> None:
        code, doc = run_tool(dense_llama_gguf, 3, 1)
        assert code == 3 and doc["ok"] is False

    def test_missing_required_owned_tensor_is_refused(self, tmp_path: Path) -> None:
        corrupted = write_dense_llama_gguf(
            tmp_path / "missing-tensor.gguf", drop=frozenset({"blk.1.attn_q.weight"})
        )
        code, doc = run_tool(corrupted, 0, 2)
        assert code == 3 and doc["ok"] is False
        assert "blk.1.attn_q.weight" in doc["error"]

    def test_whole_model_load_still_works_through_the_range_loader(
        self, dense_llama_gguf: Path
    ) -> None:
        code, doc = run_tool(dense_llama_gguf, 0, 4)
        assert code == 0
        report = parse_owned_range_report(doc)
        assert report.is_head and report.is_tail
        assert report.has_token_embeddings and report.has_output_head


def test_tool_binary_gate_points_at_the_locked_build() -> None:
    # The gate must name the deterministic lane's output, never a downloaded binary.
    assert DEFAULT_BINARY.name == "meshnet-range-report"
    assert "llama.cpp" in DEFAULT_BINARY.parts
    assert DEFAULT_BINARY.parent.name == "bin"
    assert DEFAULT_BINARY.parent.parent.name == "build"