"""US-012 tests for the real PyTorch node backend."""

from collections import OrderedDict
import json
import os
from pathlib import Path
import sys
import threading
import time
import types
import urllib.request

import pytest

from meshnet_node.model_backend import (
    InsufficientVRAMError,
    PartialModelLoadUnsupported,
    KVCacheMiss,
    TensorPayload,
    TorchModelShard,
    _call_layer,
    _checkpoint_tensor_name_for_model,
    _load_partial_model_from_snapshot,
    _should_partial_materialize_shard,
    _decoder_attention_mask,
    _int_tensor_header,
    _tensor_from_bfloat16_bytes,
    _torch_cuda_is_executable,
    build_quantization_config,
    validate_quantization,
)
from meshnet_node.torch_server import TorchNodeServer


def _require_functional_torch():
    """Skip tensor-behaviour tests when the installed torch namespace is incomplete."""
    torch = pytest.importorskip("torch")
    required = ("tensor", "zeros", "ones", "arange", "bfloat16", "long")
    if not all(hasattr(torch, name) for name in required):
        pytest.skip("requires a functional PyTorch tensor runtime")
    return torch


class _FakeBackend:
    model_id = "fake-model"
    total_layers = 12
    is_head = True
    is_tail = False

    def encode_prompt(self, prompt: str) -> TensorPayload:
        assert prompt == "The capital of France is"
        return TensorPayload(
            body=b"\x00" * (1 * 6 * 8 * 2),
            shape=[1, 6, 8],
            attention_mask_header=None,
            position_ids_header=None,
        )

    def forward_bytes(
        self,
        body,
        shape,
        attention_mask_header,
        position_ids_header,
        start_layer=None,
        **kwargs,  # noqa: ARG002
    ):
        assert shape == [1, 6, 8]
        return TensorPayload(
            body=body,
            shape=shape,
            attention_mask_header=attention_mask_header,
            position_ids_header=position_ids_header,
        )


class _FakeTailBackend(_FakeBackend):
    is_head = False
    is_tail = True

    def forward_bytes(
        self,
        body,
        shape,
        attention_mask_header,
        position_ids_header,
        start_layer=None,
        **kwargs,  # noqa: ARG002
    ):
        assert len(body) == 1 * 6 * 8 * 2
        return " Paris"


class _FakeFullBackend(_FakeBackend):
    is_head = True
    is_tail = True

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 16,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
        assert max_new_tokens == 7
        assert temperature == 1.0
        assert top_p == 1.0
        return "56"

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
        return 8

    def count_text_tokens(self, text: str) -> int:
        assert text == "56"
        return 1


class _FakeChatTokenizer:
    eos_token = ""

    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
        assert add_generation_prompt is True
        assert tokenize is False
        return "debug prompt"


class _FakePipelineHeadBackend(_FakeBackend):
    tokenizer = _FakeChatTokenizer()

    def encode_prompt(self, prompt: str) -> TensorPayload:
        assert prompt.startswith("debug prompt")
        return TensorPayload(
            body=b"\x00" * (1 * 6 * 8 * 2),
            shape=[1, 6, 8],
            attention_mask_header=None,
            position_ids_header=None,
        )


class _FakePipelineTailBackend(_FakeTailBackend):
    def __init__(self) -> None:
        self.start_layers: list[int | None] = []

    def forward_bytes(
        self,
        body,
        shape,
        attention_mask_header,
        position_ids_header,
        start_layer=None,
        **kwargs,  # noqa: ARG002
    ):
        self.start_layers.append(start_layer)
        assert len(body) == 1 * 6 * 8 * 2
        return " token"


class _BlockingStreamingTailBackend(_FakeTailBackend):
    def __init__(self, second_token_release: threading.Event) -> None:
        self._release = second_token_release
        self.calls = 0

    def forward_bytes(
        self,
        body,
        shape,
        attention_mask_header,
        position_ids_header,
        start_layer=None,
        **kwargs,  # noqa: ARG002
    ):
        self.calls += 1
        if self.calls == 1:
            return " first"
        self._release.wait(timeout=3.0)
        return " second"


def test_quantization_flag_validation():
    "Quantization flag validation\n\nTags: model, node, real-inference"
    assert validate_quantization("bfloat16") == "bfloat16"
    assert validate_quantization("int8") == "int8"
    assert validate_quantization("nf4") == "nf4"
    with pytest.raises(ValueError, match="quantization"):
        validate_quantization("float32")


def test_node_package_declares_torch_dependency():
    "Node package declares torch dependency\n\nTags: model, node, real-inference"
    pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")

    assert '"torch>=' in pyproject


def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
    "Bitsandbytes configs are created lazily\n\nTags: model, node, real-inference"
    calls = []

    class FakeBitsAndBytesConfig:
        def __init__(self, **kwargs):
            calls.append(kwargs)

    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16"))
    monkeypatch.setitem(
        sys.modules,
        "transformers",
        types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig),
    )

    assert build_quantization_config("bfloat16") is None
    build_quantization_config("int8")
    build_quantization_config("nf4")

    assert calls == [
        {"load_in_8bit": True},
        {
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_compute_dtype": "bf16",
        },
    ]


def test_rocm_inventory_without_executable_kernels_is_not_used_as_cuda():
    "Rocm inventory without executable kernels is not used as cuda\n\nTags: model, node, real-inference"
    class FakeCuda:
        @staticmethod
        def is_available():
            return True

        @staticmethod
        def synchronize():
            raise AssertionError("synchronize should not run after empty() fails")

    fake_torch = types.SimpleNamespace(
        cuda=FakeCuda(),
        empty=lambda *args, **kwargs: (_ for _ in ()).throw(
            RuntimeError("HIP error: invalid device function")
        ),
    )

    assert _torch_cuda_is_executable(fake_torch) is False


def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
    "Head forward accepts text prompt and returns bfloat16 activations\n\nTags: model, node, real-inference"
    node = TorchNodeServer(backend=_FakeBackend())
    port = node.start()
    try:
        payload = json.dumps({"prompt": "The capital of France is"}).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/forward",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = resp.read()
            headers = {key.lower(): value for key, value in resp.headers.items()}

        assert len(body) == 1 * 6 * 8 * 2
        assert headers["x-meshnet-shape"] == "1,6,8"
        assert headers["x-meshnet-dtype"] == "bfloat16"
        assert headers["x-meshnet-wire"] == "2"
    finally:
        node.stop()


def test_tail_forward_returns_text_completion_from_binary_activations():
    "Tail forward returns text completion from binary activations\n\nTags: model, node, real-inference"
    node = TorchNodeServer(backend=_FakeTailBackend())
    port = node.start()
    try:
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/forward",
            data=b"\x00" * (1 * 6 * 8 * 2),
            headers={
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Shape": "1,6,8",
                "X-Meshnet-Dtype": "bfloat16",
                "X-Meshnet-Session": "session-1",
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": "1",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())

        assert body == {"text": " Paris"}
        assert node.received_activations
        assert node.forward_chunk_count == 1
    finally:
        node.stop()


def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys):
    "Full model chat completion uses generation not single token decode\n\nTags: model, node, real-inference"
    node = TorchNodeServer(backend=_FakeFullBackend())
    port = node.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "What is 7 times 8?"}],
            "max_tokens": 7,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Request-Id": "req-test-123",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())

        assert body["choices"][0]["message"]["content"] == "56"
        assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
    finally:
        node.stop()

    out = capsys.readouterr().out
    assert "  [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out
    assert "  [node] chat complete tokens=1 elapsed_s=" in out


def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
    "Pipeline hop logs are suppressed without debug\n\nTags: model, node, real-inference"
    tail_backend = _FakePipelineTailBackend()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=tail_backend)
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert body["choices"][0]["message"]["content"] == " token"
    assert tail_backend.start_layers == [22]
    assert "pipeline hop 0:" not in out
    assert "pipeline hop 0 returned text" not in out


def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
    "Pipeline hop logs are enabled with debug\n\nTags: model, node, real-inference"
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert f"  [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
    assert "  [node] pipeline hop 0 returned text=' token'" in out


def test_split_shard_chat_streams_each_generated_token_incrementally():
    "Split shard chat streams each generated token incrementally\n\nTags: model, node, real-inference, streaming"
    release_second = threading.Event()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
    head_port = head.start()
    tail_port = tail.start()
    response = None
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "stream": True,
            "max_tokens": 2,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        response = urllib.request.urlopen(req, timeout=5)

        first_token_line = ""
        deadline = time.time() + 2.0
        while time.time() < deadline:
            line = response.readline().decode()
            if '"content": " first"' in line:
                first_token_line = line
                break

        assert first_token_line
        assert not release_second.is_set()
        release_second.set()
        rest = response.read().decode()
    finally:
        release_second.set()
        if response is not None:
            response.close()
        head.stop()
        tail.stop()

    assert '"content": " second"' in rest
    assert "data: [DONE]" in rest


def test_current_requests_snapshot_while_generating():
    "Current requests snapshot while generating\n\nTags: model, node, real-inference"
    release_second = threading.Event()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
    head_port = head.start()
    tail_port = tail.start()
    response = None
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "stream": True,
            "max_tokens": 2,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Request-Id": "req-live-1",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        response = urllib.request.urlopen(req, timeout=5)
        deadline = time.time() + 2.0
        while time.time() < deadline:
            live = head.current_requests
            if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1:
                break
            time.sleep(0.02)
        assert head.current_requests
        snap = head.current_requests[0]
        assert snap["request_id"] == "req-live-1"
        assert snap["tokens"] >= 1
        assert snap["tokens_per_sec"] >= 0
        assert snap["routing_complete"] is True
        release_second.set()
        response.read()
    finally:
        release_second.set()
        if response is not None:
            response.close()
        head.stop()
        tail.stop()

    assert head.current_requests == []


def test_distributed_generating_log_includes_tps(capsys):
    "Distributed generating log includes tps\n\nTags: model, node, real-inference"
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert "generating step=1/1" in out
    assert " tps=" in out
    assert "generation complete session=" in out
    assert "tokens=1" in out
    assert out.count("generating step=1/1") == 1


def test_int_tensor_header_serializes_torch_tensors():
    "Int tensor header serializes torch tensors\n\nTags: model, node, real-inference"
    torch = _require_functional_torch()

    header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))

    assert header.startswith("1,3:")


def test_bfloat16_wire_decode_views_owned_bytes_without_float32_round_trip():
    """Activation decode stays bf16 and does not clone bytes into bytearray.

    Tags: model, performance, wire
    """
    torch = _require_functional_torch()
    body = torch.tensor([[1, 2]], dtype=torch.bfloat16).view(torch.uint8).numpy().tobytes()

    decoded = _tensor_from_bfloat16_bytes(body, [1, 2], torch)

    assert decoded.dtype == torch.bfloat16
    assert decoded.tolist() == [[1.0, 2.0]]


def test_decoder_attention_mask_is_causal_float_mask():
    "Decoder attention mask is causal float mask\n\nTags: model, node, real-inference"
    torch = _require_functional_torch()

    hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
    mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)

    assert mask.shape == (1, 1, 3, 3)
    assert mask.dtype == torch.bfloat16
    assert mask[0, 0, 0, 1] < 0
    assert mask[0, 0, 2, 0] == 0


def test_call_layer_passes_rotary_position_embeddings():
    "Call layer passes rotary position embeddings\n\nTags: model, node, real-inference"
    class NeedsPositionEmbeddings:
        def __call__(self, hidden_states, **kwargs):
            assert kwargs["position_embeddings"] == "rotary"
            return hidden_states

    hidden = _call_layer(
        NeedsPositionEmbeddings(),
        "hidden",
        attention_mask=None,
        position_ids="positions",
        position_embeddings="rotary",
    )

    assert hidden == "hidden"


def _fake_cache_shard(torch, *, max_sessions=16, ttl=600.0):
    class RecordingLayer:
        def __init__(self, index):
            self.index = index
            self.calls = []

        def __call__(self, hidden_states, **kwargs):
            self.calls.append({
                "shape": tuple(hidden_states.shape),
                "use_cache": kwargs.get("use_cache"),
                "past_key_value": kwargs.get("past_key_value"),
            })
            present = {
                "layer": self.index,
                "shape": tuple(hidden_states.shape),
                "opaque": object(),
            }
            return hidden_states + (self.index + 1), present

    shard = object.__new__(TorchModelShard)
    shard.shard_start = 0
    shard.shard_end = 1
    shard.torch = torch
    shard.model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[]))
    shard.layers = [RecordingLayer(0), RecordingLayer(1)]
    shard._session_cache = OrderedDict()
    shard._cache_max_sessions = max_sessions
    shard._cache_ttl_seconds = ttl
    return shard


def test_shard_cache_prefill_then_decode_reuses_opaque_layer_state():
    "Shard cache prefill then decode reuses opaque layer state\n\nTags: cache, model, node, real-inference"
    torch = _require_functional_torch()
    shard = _fake_cache_shard(torch)

    prefill_hidden = torch.zeros((1, 4, 2), dtype=torch.bfloat16)
    prefill_mask = torch.ones((1, 4), dtype=torch.long)
    prefill_positions = torch.arange(4, dtype=torch.long).reshape(1, 4)
    shard._run_layers(
        prefill_hidden,
        prefill_mask,
        prefill_positions,
        session_id="session-1",
        cache_mode="prefill",
        seq_len=4,
    )

    assert len(shard._session_cache) == 1
    cached_states = next(iter(shard._session_cache.values())).layer_states
    assert len(cached_states) == 2
    assert cached_states[0]["shape"] == (1, 4, 2)

    decode_hidden = torch.zeros((1, 1, 2), dtype=torch.bfloat16)
    decode_mask = torch.ones((1, 5), dtype=torch.long)
    decode_positions = torch.tensor([[4]], dtype=torch.long)
    shard._run_layers(
        decode_hidden,
        decode_mask,
        decode_positions,
        session_id="session-1",
        cache_mode="decode",
        seq_len=5,
    )

    assert shard.layers[0].calls[-1]["shape"] == (1, 1, 2)
    assert shard.layers[0].calls[-1]["past_key_value"] is cached_states[0]
    assert shard.layers[1].calls[-1]["past_key_value"] is cached_states[1]
    assert next(iter(shard._session_cache.values())).seq_len == 5


def test_shard_cache_decode_miss_is_explicit():
    "Shard cache decode miss is explicit\n\nTags: cache, model, node, real-inference"
    torch = _require_functional_torch()
    shard = _fake_cache_shard(torch)

    with pytest.raises(KVCacheMiss):
        shard._run_layers(
            torch.zeros((1, 1, 2), dtype=torch.bfloat16),
            torch.ones((1, 5), dtype=torch.long),
            torch.tensor([[4]], dtype=torch.long),
            session_id="missing",
            cache_mode="decode",
            seq_len=5,
        )


def test_shard_cache_lru_bounds_sessions():
    "Shard cache lru bounds sessions\n\nTags: cache, model, node, real-inference"
    torch = _require_functional_torch()
    shard = _fake_cache_shard(torch, max_sessions=1)

    for session in ("old", "new"):
        shard._run_layers(
            torch.zeros((1, 2, 2), dtype=torch.bfloat16),
            torch.ones((1, 2), dtype=torch.long),
            torch.arange(2, dtype=torch.long).reshape(1, 2),
            session_id=session,
            cache_mode="prefill",
            seq_len=2,
        )

    assert list(shard._session_cache.keys()) == [("new", 0, 1)]


def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):
    "Partial materialize guard requires local non full non quantized snapshot\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')

    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is True
    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        0,
        39,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is True
    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=True,
    ) is False
    assert _should_partial_materialize_shard(
        "repo/model",
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is False


def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm():
    "Checkpoint tensor name remapped for text only causal lm\n\nTags: model, node, real-inference"
    class TextOnlyModel:
        def __init__(self):
            self.model = types.SimpleNamespace(layers=[])

    model = TextOnlyModel()
    assert _checkpoint_tensor_name_for_model(
        model,
        "model.language_model.layers.0.mlp.gate.weight",
    ) == "model.layers.0.mlp.gate.weight"
    assert _checkpoint_tensor_name_for_model(
        model,
        "model.language_model.embed_tokens.weight",
    ) == "model.embed_tokens.weight"


def test_checkpoint_tensor_name_kept_for_multimodal_backbone():
    "Checkpoint tensor name kept for multimodal backbone\n\nTags: model, node, real-inference"
    class MultimodalModel:
        def __init__(self):
            self.model = types.SimpleNamespace(language_model=types.SimpleNamespace())

    model = MultimodalModel()
    name = "model.language_model.layers.0.mlp.gate.weight"
    assert _checkpoint_tensor_name_for_model(model, name) == name


def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path):
    "Partial snapshot loader remaps language model checkpoint keys\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text(json.dumps({
        "text_config": {"num_hidden_layers": 3},
    }))
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
        }
    }))
    (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")

    class FakeModule:
        def __init__(self):
            self.to_calls = []

        def to(self, device):
            self.to_calls.append(device)
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[FakeModule(), FakeModule(), FakeModule()],
                rotary_emb=FakeModule(),
            )

        def tie_weights(self):
            pass

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace(
                text_config=types.SimpleNamespace(num_hidden_layers=3),
                get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
            )

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            return FakeModel()

    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append(tensor_name)

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            self.filename = Path(filename).name

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensor_name

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=UnusedContext,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]


def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path):
    # Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.*
    # tensors that the text-only CausalLM never builds — they must be skipped,
    # not assigned (assignment raises AttributeError: 'mtp' / 'visual').
    "Partial snapshot loader skips tensors absent from causal lm\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text(json.dumps({
        "text_config": {"num_hidden_layers": 3},
    }))
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
            "mtp.layers.1.input_layernorm.weight": "shard-2.safetensors",
            "model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors",
        }
    }))
    (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")

    class FakeModule:
        def to(self, device):
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[FakeModule(), FakeModule(), FakeModule()],
                rotary_emb=FakeModule(),
            )

        def tie_weights(self):
            pass

        def state_dict(self):
            return {"model.layers.1.self_attn.q_proj.weight": None}

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace(
                text_config=types.SimpleNamespace(num_hidden_layers=3),
                get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
            )

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            return FakeModel()

    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append(tensor_name)

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            pass

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensor_name

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=UnusedContext,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]


def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
    "Partial snapshot loader materializes only assigned tensors\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.embed_tokens.weight": "shard-1.safetensors",
            "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors",
            "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
            "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors",
            "model.norm.weight": "shard-3.safetensors",
            "lm_head.weight": "shard-3.safetensors",
        }
    }))
    for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"):
        (snapshot_dir / rel).write_bytes(b"stub")

    class FakeModule:
        def __init__(self, name):
            self.name = name
            self.to_calls = []

        def to(self, device):
            self.to_calls.append(device)
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                embed_tokens=FakeModule("embed"),
                layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")],
                rotary_emb=FakeModule("rotary"),
                norm=FakeModule("norm"),
            )
            self.lm_head = FakeModule("lm_head")
            self.tie_weights_called = 0

        def tie_weights(self):
            self.tie_weights_called += 1

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            assert model_id == str(snapshot_dir)
            return types.SimpleNamespace(num_hidden_layers=3)

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            assert cfg.num_hidden_layers == 3
            assert torch_dtype == "bf16"
            return FakeModel()

    class EmptyWeights:
        def __init__(self):
            self.entered = 0
            self.exited = 0

        def __call__(self):
            return self

        def __enter__(self):
            self.entered += 1
            return None

        def __exit__(self, exc_type, exc, tb):
            self.exited += 1
            return False

    init_empty_weights = EmptyWeights()
    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append((tensor_name, device, value, dtype))

    tensors = {
        "shard-1.safetensors": {
            "model.embed_tokens.weight": "embed",
            "model.layers.0.self_attn.q_proj.weight": "layer0",
        },
        "shard-2.safetensors": {
            "model.layers.1.self_attn.q_proj.weight": "layer1",
        },
        "shard-3.safetensors": {
            "model.layers.2.self_attn.q_proj.weight": "layer2",
            "model.norm.weight": "norm",
            "lm_head.weight": "lm_head",
        },
    }

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            assert framework == "pt"
            assert device == "cpu"
            self.filename = Path(filename).name

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensors[self.filename][tensor_name]

    model = _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=init_empty_weights,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert init_empty_weights.entered == 1
    assert init_empty_weights.exited == 1
    assert model.tie_weights_called == 1
    assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"]
    assert model.model.layers[1].to_calls == ["cpu:0"]
    assert model.model.layers[0].to_calls == []
    assert model.model.layers[2].to_calls == []
    assert model.model.embed_tokens.to_calls == []
    assert model.model.norm.to_calls == []
    assert model.lm_head.to_calls == []
    assert model.model.rotary_emb.to_calls == ["cpu:0"]


def _build_partial_snapshot_fixture(snapshot_dir: Path, *, num_layers: int = 4) -> dict[str, dict[str, int]]:
    """Minimal HF snapshot with per-tensor numel metadata for memory-scaling tests."""
    weight_map: dict[str, str] = {
        "model.embed_tokens.weight": "shard-head.safetensors",
        "model.norm.weight": "shard-tail.safetensors",
        "lm_head.weight": "shard-tail.safetensors",
    }
    tensor_numel: dict[str, dict[str, int]] = {
        "shard-head.safetensors": {"model.embed_tokens.weight": 10_000},
        "shard-tail.safetensors": {
            "model.norm.weight": 512,
            "lm_head.weight": 10_000,
        },
    }
    for layer in range(num_layers):
        rel = f"shard-layer-{layer}.safetensors"
        weight_map[f"model.layers.{layer}.self_attn.q_proj.weight"] = rel
        weight_map[f"model.layers.{layer}.mlp.down_proj.weight"] = rel
        per_layer = 1_000 * (layer + 1)
        tensor_numel[rel] = {
            f"model.layers.{layer}.self_attn.q_proj.weight": per_layer,
            f"model.layers.{layer}.mlp.down_proj.weight": per_layer,
        }
        (snapshot_dir / rel).write_bytes(b"stub")

    (snapshot_dir / "config.json").write_text(json.dumps({"num_hidden_layers": num_layers}))
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))
    (snapshot_dir / "shard-head.safetensors").write_bytes(b"stub")
    (snapshot_dir / "shard-tail.safetensors").write_bytes(b"stub")
    return tensor_numel


def _partial_load_materialized_numel(
    snapshot_dir: Path,
    shard_start: int,
    shard_end: int,
    tensor_numel: dict[str, dict[str, int]],
) -> int:
    """Return the summed numel of checkpoint tensors assigned to one shard load."""

    class FakeModule:
        def to(self, device):
            return self

    class FakeModel:
        def __init__(self, num_layers: int):
            self.model = types.SimpleNamespace(
                embed_tokens=FakeModule(),
                layers=[FakeModule() for _ in range(num_layers)],
                rotary_emb=FakeModule(),
                norm=FakeModule(),
            )
            self.lm_head = FakeModule()

        def tie_weights(self):
            pass

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            assert model_id == str(snapshot_dir)
            return types.SimpleNamespace(num_hidden_layers=num_layers)

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            return FakeModel(cfg.num_hidden_layers)

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    num_layers = json.loads((snapshot_dir / "config.json").read_text())["num_hidden_layers"]
    loaded_numel = 0

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        nonlocal loaded_numel
        loaded_numel += int(getattr(value, "numel", lambda: 0)())

    class FakeTensor:
        def __init__(self, numel: int):
            self._numel = numel

        def numel(self) -> int:
            return self._numel

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            self.filename = Path(filename).name

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return FakeTensor(tensor_numel[self.filename][tensor_name])

    _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        shard_start,
        shard_end,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=UnusedContext,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )
    return loaded_numel


def test_partial_snapshot_resident_weight_numel_scales_with_shard(tmp_path):
    "Partial load materializes only assigned-layer weights, not the full checkpoint.\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    tensor_numel = _build_partial_snapshot_fixture(snapshot_dir, num_layers=4)

    middle_numel = _partial_load_materialized_numel(snapshot_dir, 1, 1, tensor_numel)
    full_numel = _partial_load_materialized_numel(snapshot_dir, 0, 3, tensor_numel)

    # Layer 1 only: two tensors at 2000 numel each.
    assert middle_numel == 4_000
    # Head embed + four layers (two tensors each, increasing sizes) + tail norm/lm_head.
    assert full_numel == 10_000 + 512 + 10_000 + 2_000 + 4_000 + 6_000 + 8_000
    assert middle_numel < full_numel // 4


@pytest.mark.integration
def test_partial_snapshot_materialized_param_count_with_real_torch(tmp_path):
    "When torch is installed, partial load leaves unassigned params on meta device.\n\nTags: model, node, real-inference"
    torch = _require_functional_torch()
    pytest.importorskip("transformers")
    pytest.importorskip("safetensors")
    pytest.importorskip("accelerate")
    from safetensors.torch import save_file
    from transformers import AutoConfig, AutoModelForCausalLM, GPT2Config

    n_layer = 4
    n_embd = 16
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    config = GPT2Config(n_layer=n_layer, n_embd=n_embd, n_head=2, n_positions=32)
    (snapshot_dir / "config.json").write_text(config.to_json_string())

    weight_map: dict[str, str] = {}
    for layer in range(n_layer):
        key = f"transformer.h.{layer}.attn.c_attn.weight"
        rel = f"layer-{layer}.safetensors"
        weight_map[key] = rel
        save_file({key: torch.ones(n_embd, 3 * n_embd)}, snapshot_dir / rel)
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))

    def count_materialized(model) -> tuple[int, int]:
        materialized = 0
        meta = 0
        for param in model.parameters():
            if param.device.type == "meta":
                meta += param.numel()
            else:
                materialized += param.numel()
        return materialized, meta

    middle = _load_partial_model_from_snapshot(
        AutoConfig,
        AutoModelForCausalLM,
        torch,
        str(snapshot_dir),
        1,
        1,
        torch.float32,
        torch.device("cpu"),
    )
    mid_mat, mid_meta = count_materialized(middle)
    assert mid_mat == n_embd * 3 * n_embd
    assert mid_meta > mid_mat * 10


def test_partial_snapshot_loader_requires_known_layer_count(tmp_path):
    "Partial snapshot loader requires known layer count\n\nTags: model, node, real-inference"
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"}
    }))
    (snapshot_dir / "shard.safetensors").write_bytes(b"stub")

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace()

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            raise AssertionError("from_config should not run without a known layer count")

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"):
        _load_partial_model_from_snapshot(
            AutoConfigStub,
            AutoModelStub,
            types.SimpleNamespace(),
            str(snapshot_dir),
            0,
            0,
            "bf16",
            "cpu:0",
            init_empty_weights_fn=lambda: UnusedContext(),
            set_tensor_fn=lambda *args, **kwargs: None,
            safe_open_fn=lambda *args, **kwargs: None,
        )


def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch):
    "Torch model shard prefers partial loader for local snapshot\n\nTags: model, node, real-inference"
    import meshnet_node.model_backend as backend

    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[object(), object(), object()],
                embed_tokens=object(),
            )
            self.config = types.SimpleNamespace(hidden_size=8)
            self.eval_called = 0

        def eval(self):
            self.eval_called += 1

    fake_model = FakeModel()
    partial_calls = []

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16"))

    class AutoModelStub:
        @staticmethod
        def from_pretrained(*args, **kwargs):
            raise AssertionError("full model load should not run for partial local shards")

    class AutoTokenizerStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            assert model_id == str(snapshot_dir)
            return types.SimpleNamespace()

    monkeypatch.setitem(
        sys.modules,
        "torch",
        types.SimpleNamespace(
            cuda=types.SimpleNamespace(is_available=lambda: False),
            device=lambda value: types.SimpleNamespace(type=value),
            bfloat16="bf16",
        ),
    )
    monkeypatch.setitem(
        sys.modules,
        "transformers",
        types.SimpleNamespace(
            AutoConfig=AutoConfigStub,
            AutoModelForCausalLM=AutoModelStub,
            AutoTokenizer=AutoTokenizerStub,
        ),
    )
    monkeypatch.setattr(
        backend,
        "_load_partial_model_from_snapshot",
        lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model,
    )

    shard = TorchModelShard(
        "repo/model",
        1,
        1,
        quantization="auto",
        cache_dir=snapshot_dir,
    )

    assert len(partial_calls) == 1
    assert shard.model is fake_model
    assert fake_model.eval_called == 1
    assert shard.total_layers == 3
    assert shard.is_head is False
    assert shard.is_tail is False


@pytest.mark.integration
def test_two_node_gpt2_completion_is_deterministic():
    "Two node gpt2 completion is deterministic\n\nTags: model, node, real-inference"
    if os.environ.get("CI"):
        pytest.skip("GPT-2 integration test is skipped in CI")
    torch = _require_functional_torch()
    pytest.importorskip("transformers")
    pytest.importorskip("safetensors")
    pytest.importorskip("accelerate")
    pytest.importorskip("bitsandbytes")
    if not torch.cuda.is_available():
        pytest.skip("GPT-2 integration test requires a CUDA GPU")

    head = TorchNodeServer(
        model_id="openai-community/gpt2",
        shard_start=0,
        shard_end=6,
        quantization="bfloat16",
    )
    tail = TorchNodeServer(
        model_id="openai-community/gpt2",
        shard_start=6,
        shard_end=12,
        quantization="bfloat16",
    )
    head_port = head.start()
    tail_port = tail.start()
    try:
        prompt_req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/forward",
            data=json.dumps({"prompt": "The capital of France is"}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(prompt_req, timeout=60) as resp:
            activation = resp.read()
            head_headers = resp.headers

        tail_req = urllib.request.Request(
            f"http://127.0.0.1:{tail_port}/forward",
            data=activation,
            headers={
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Shape": head_headers["X-Meshnet-Shape"],
                "X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"],
                "X-Meshnet-Session": "gpt2-session",
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": "1",
                "X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"],
                "X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"],
            },
            method="POST",
        )
        with urllib.request.urlopen(tail_req, timeout=60) as resp:
            body = json.loads(resp.read())

        assert body["text"].strip()
        assert body["text"] == " Paris"
    finally:
        head.stop()
        tail.stop()