"""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,
    ShardCacheMiss,
    TensorPayload,
    TorchModelShard,
    _call_layer,
    _checkpoint_tensor_name_for_model,
    _load_partial_model_from_snapshot,
    _should_partial_materialize_shard,
    _decoder_attention_mask,
    _int_tensor_header,
    _torch_cuda_is_executable,
    build_quantization_config,
    validate_quantization,
)
from meshnet_node.torch_server import TorchNodeServer


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():
    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():
    pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")

    assert '"torch>=' in pyproject


def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
    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():
    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():
    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():
    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):
    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):
    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):
    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():
    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():
    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):
    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 tokens=1" in out
    assert out.count("generating step=1/1") == 1


def test_int_tensor_header_serializes_torch_tensors():
    torch = pytest.importorskip("torch")

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

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


def test_decoder_attention_mask_is_causal_float_mask():
    torch = pytest.importorskip("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():
    class NeedsPositionEmbeddings:
        def __call__(self, hidden_states, **kwargs):
            assert kwargs["position_embeddings"] == "rotary"
            return hidden_states

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

    assert hidden == "hidden"
    assert cache_state is None


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():
    torch = pytest.importorskip("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():
    torch = pytest.importorskip("torch")
    shard = _fake_cache_shard(torch)

    with pytest.raises(ShardCacheMiss):
        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():
    torch = pytest.importorskip("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):
    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():
    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():
    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):
    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').
    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):
    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 test_partial_snapshot_loader_requires_known_layer_count(tmp_path):
    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):
    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: 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():
    if os.environ.get("CI"):
        pytest.skip("GPT-2 integration test is skipped in CI")
    torch = pytest.importorskip("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()