neuron-tai/tests/test_node_startup.py

"""US-004 integration tests: node self-configuring startup sequence."""

import json
import io
import os
import sys
import tarfile
import time
import types
import urllib.request
from pathlib import Path

import pytest

from meshnet_node.downloader import download_shard, write_shard_archive
from meshnet_node.hardware import detect_hardware, benchmark_throughput
from meshnet_node.startup import (
    _configure_torch_threads,
    _hardware_label,
    _infer_relay_url_from_tracker,
    _memory_budget,
    _probationary_status_line,
    run_startup,
)
from meshnet_node.wallet import _b58encode, load_or_create_wallet
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.server import TrackerServer


# ---------------------------------------------------------------------------
# Unit tests — hardware, wallet, downloader
# ---------------------------------------------------------------------------


def test_detect_hardware_returns_valid_profile():
    """Hardware detection always returns a dict with required keys."""
    hw = detect_hardware()
    assert hw["device"] in {"cuda", "cpu"}
    assert isinstance(hw.get("vram_mb"), int)
    assert isinstance(hw.get("ram_mb"), int)
    assert hw["ram_mb"] > 0
    if hw["device"] == "cpu":
        assert hw["gpu_name"] is None or isinstance(hw["gpu_name"], str)
        assert hw["vram_mb"] >= 0
    else:
        assert isinstance(hw["gpu_name"], str) and hw["gpu_name"]
        assert hw["vram_mb"] > 0


def test_windows_ram_fallback_is_used_when_sysconf_is_unavailable(monkeypatch):
    """Windows hosts do not have os.sysconf; RAM must not collapse to 0 MB."""
    import meshnet_node.hardware as hardware_mod

    monkeypatch.setattr(
        hardware_mod.os,
        "sysconf",
        lambda _name: (_ for _ in ()).throw(AttributeError()),
        raising=False,
    )
    monkeypatch.setattr(hardware_mod, "_detect_windows_ram_mb", lambda: 64 * 1024)

    assert hardware_mod._detect_ram_mb() == 64 * 1024


def test_windows_gpu_memory_fallback_preserves_cpu_execution(monkeypatch):
    """A Windows-visible GPU is reported, but CUDA execution is not claimed without CUDA."""
    import meshnet_node.hardware as hardware_mod

    calls = []

    class FakeResult:
        def __init__(self, stdout):
            self.returncode = 0
            self.stdout = stdout

    def fake_run(command, *args, **kwargs):
        calls.append(command)
        joined = " ".join(command)
        if "nvidia-smi" in joined:
            raise FileNotFoundError
        if "Win32_ComputerSystem" in joined:
            return FakeResult(str(80 * 1024 * 1024 * 1024))
        if "Win32_VideoController" in joined:
            return FakeResult('{"Name":"NVIDIA GeForce RTX Laptop GPU","AdapterRAM":8589934592}')
        raise AssertionError(command)

    monkeypatch.setattr(
        hardware_mod.os,
        "sysconf",
        lambda _name: (_ for _ in ()).throw(AttributeError()),
        raising=False,
    )
    monkeypatch.setattr(hardware_mod.subprocess, "run", fake_run)
    monkeypatch.setattr(hardware_mod, "_detect_windows_ram_mb", lambda: 80 * 1024)
    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(cuda=types.SimpleNamespace(is_available=lambda: False)))

    hw = hardware_mod.detect_hardware()

    assert hw["device"] == "cpu"
    assert hw["gpu_name"] == "NVIDIA GeForce RTX Laptop GPU"
    assert hw["vram_mb"] == 8192
    assert hw["shared_vram_mb"] == 40 * 1024
    assert hw["ram_mb"] == 80 * 1024


def test_nvidia_smi_without_torch_cuda_keeps_cpu_execution(monkeypatch):
    """nvidia-smi proves GPU inventory, not that this Python can execute CUDA."""
    import meshnet_node.hardware as hardware_mod

    class FakeResult:
        returncode = 0
        stdout = "NVIDIA GeForce RTX 4060 Laptop GPU, 8188\n"

    fake_torch = types.SimpleNamespace(cuda=types.SimpleNamespace(is_available=lambda: False))

    monkeypatch.setattr(hardware_mod, "_detect_ram_mb", lambda: 80 * 1024)
    monkeypatch.setattr(hardware_mod.subprocess, "run", lambda *a, **kw: FakeResult())
    monkeypatch.setitem(sys.modules, "torch", fake_torch)

    hw = hardware_mod.detect_hardware()

    assert hw["device"] == "cpu"
    assert hw["cuda_available"] is False
    assert hw["gpu_name"] == "NVIDIA GeForce RTX 4060 Laptop GPU"
    assert hw["vram_mb"] == 8188
    assert hw["ram_mb"] == 80 * 1024


def test_memory_budget_uses_ram_for_cpu_and_shared_memory_for_cuda():
    assert _memory_budget("cpu", vram_mb=8192, ram_mb=80 * 1024, shared_vram_mb=40 * 1024) == (
        80 * 1024,
        "RAM",
    )
    assert _memory_budget("cuda", vram_mb=8192, ram_mb=80 * 1024, shared_vram_mb=40 * 1024) == (
        48 * 1024,
        "VRAM + shared RAM",
    )


def test_hardware_label_marks_inventory_only_gpu_as_cuda_inactive():
    assert _hardware_label("cpu", "NVIDIA GeForce RTX 4060 Laptop GPU") == "CPU (CUDA inactive)"
    assert _hardware_label("cpu", None) == "CPU"
    assert _hardware_label("cuda", "NVIDIA GeForce RTX 4060 Laptop GPU") == "CUDA"


def test_benchmark_throughput_cpu_returns_positive():
    """CPU benchmark returns a positive float greater than the 1.0 error fallback."""
    result = benchmark_throughput("cpu")
    assert isinstance(result, float)
    assert result > 1.0, f"expected benchmark > 1.0, got {result}"


def test_benchmark_throughput_fallback_on_bad_device():
    """benchmark_throughput returns 1.0 (not raises) when device is invalid."""
    result = benchmark_throughput("invalid_device_xyz")
    assert result == 1.0


def test_configure_torch_threads_applies_explicit_settings(monkeypatch):
    """Node startup can tune PyTorch CPU thread pools before loading a model."""
    calls: dict[str, int] = {}

    fake_torch = types.SimpleNamespace(
        set_num_threads=lambda value: calls.update({"threads": value}),
        set_num_interop_threads=lambda value: calls.update({"interop_threads": value}),
        get_num_threads=lambda: calls["threads"],
        get_num_interop_threads=lambda: calls["interop_threads"],
    )

    monkeypatch.setitem(sys.modules, "torch", fake_torch)
    monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
    monkeypatch.delenv("MKL_NUM_THREADS", raising=False)

    active = _configure_torch_threads(torch_threads=12, torch_interop_threads=2)

    assert calls == {"threads": 12, "interop_threads": 2}
    assert os.environ["OMP_NUM_THREADS"] == "12"
    assert os.environ["MKL_NUM_THREADS"] == "12"
    assert active == {"torch_threads": 12, "torch_interop_threads": 2}


def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
    """benchmark_tokens_per_sec from the benchmark is included in the tracker registration."""
    import meshnet_node.startup as startup_mod

    captured: dict = {}
    thread_calls: dict[str, int] = {}

    class FakeNode:
        backend = None
        tracker_node_id = None

        def start(self):
            return 7099

        def stop(self):
            pass

        def apply_tracker_directives(self, directives):
            return None

    monkeypatch.setattr(startup_mod, "detect_hardware",
                        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384})
    monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (42.5, True, None))
    monkeypatch.setitem(
        sys.modules,
        "torch",
        types.SimpleNamespace(
            set_num_threads=lambda value: thread_calls.update({"threads": value}),
            set_num_interop_threads=lambda value: thread_calls.update({"interop_threads": value}),
            get_num_threads=lambda: thread_calls["threads"],
            get_num_interop_threads=lambda: thread_calls["interop_threads"],
        ),
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode())
    monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
    monkeypatch.setattr(startup_mod, "RelayHttpBridge", None)
    monkeypatch.setattr(startup_mod, "_get_json",
                        lambda _url, timeout=10.0: {"relay_url": None, "nodes": []})
    monkeypatch.setattr(startup_mod, "_post_json",
                        lambda _url, payload, timeout=10.0: (captured.update(payload) or {"node_id": "x"}))
    monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)

    node = run_startup(
        tracker_url="http://localhost:8080",
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        shard_start=0,
        shard_end=23,
        wallet_path=tmp_path / "wallet.json",
        torch_threads=8,
        torch_interop_threads=1,
    )
    node.stop()

    assert captured.get("benchmark_tokens_per_sec") == 42.5
    assert captured["hardware_profile"]["torch_threads"] == 8
    assert captured["hardware_profile"]["torch_interop_threads"] == 1
    assert captured["hardware_profile"]["benchmark_device"] == "cpu"
    assert captured["hardware_profile"]["benchmark_ok"] is True


def test_real_model_startup_passes_download_dir_and_kimi_metadata(monkeypatch, tmp_path):
    import meshnet_node.startup as startup_mod

    captured_registration: dict = {}
    captured_torch_kwargs: dict = {}

    class FakeBackend:
        total_layers = 61

    class FakeNode:
        backend = FakeBackend()

        def __init__(self, **kwargs):
            captured_torch_kwargs.update(kwargs)

        def start(self):
            return 7099

        def stop(self):
            pass

        def apply_tracker_directives(self, directives):
            return None

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
    )
    monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (42.5, True, None))
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeNode)
    monkeypatch.setattr(startup_mod, "load_or_create_wallet", lambda **_kw: (b"", b"", "wallet-kimi"))
    monkeypatch.setattr(startup_mod, "_get_json", lambda _url, timeout=10.0: {"relay_url": None, "nodes": []})
    monkeypatch.setattr(
        startup_mod,
        "_post_json",
        lambda _url, payload, timeout=10.0: (
            captured_registration.update(payload) or {"node_id": "node-kimi"}
        ),
    )
    monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)

    cache_dir = tmp_path / "models"
    node = run_startup(
        tracker_url="http://localhost:8080",
        model_id="unsloth/Kimi-K2.7-Code",
        shard_start=0,
        shard_end=60,
        wallet_path=tmp_path / "wallet.json",
        cache_dir=cache_dir,
    )
    node.stop()

    assert captured_torch_kwargs["cache_dir"] == cache_dir
    assert captured_registration["model_metadata"]["total_parameters"] == "1T"
    assert captured_registration["model_metadata"]["activated_parameters"] == "32B"
    assert captured_registration["model_metadata"]["context_length"] == 256000


def test_cuda_benchmark_failure_is_registered_for_inventory_only_gpu(monkeypatch, tmp_path, capsys):
    import meshnet_node.startup as startup_mod

    captured: dict = {}

    class FakeNode:
        backend = None

        def start(self):
            return 7099

        def stop(self):
            pass

    def fake_benchmark(device):
        if device == "cuda":
            return 1.0, False, "AssertionError: Torch not compiled with CUDA enabled"
        return 55.0, True, None

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {
            "device": "cpu",
            "gpu_name": "NVIDIA GeForce RTX 4060 Laptop GPU",
            "vram_mb": 8188,
            "dedicated_vram_mb": 8188,
            "shared_vram_mb": 40555,
            "ram_mb": 81111,
            "cuda_available": False,
        },
    )
    monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", fake_benchmark)
    monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode())
    monkeypatch.setattr(startup_mod, "_post_json",
                        lambda _url, payload, timeout=10.0: (captured.update(payload) or {"node_id": "x"}))
    monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)

    node = run_startup(
        tracker_url="http://localhost:8080",
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        shard_start=0,
        shard_end=23,
        wallet_path=tmp_path / "wallet.json",
    )
    node.stop()

    output = capsys.readouterr().out
    assert "CUDA benchmark unavailable" in output
    assert "Hardware:     CPU (CUDA inactive)" in output
    hw = captured["hardware_profile"]
    assert hw["cuda_benchmark_ok"] is False
    assert "Torch not compiled with CUDA enabled" in hw["cuda_benchmark_error"]
    assert hw["benchmark_device"] == "cpu"
    assert hw["benchmark_ok"] is True
    assert captured["ram_bytes"] == 81111 * 1024 * 1024
    assert captured["vram_bytes"] == 0


def test_wallet_generates_new_keypair(tmp_path):
    """A new wallet is created when none exists, saved to disk."""
    wallet_file = tmp_path / "wallet.json"
    assert not wallet_file.exists()

    secret, public, address = load_or_create_wallet(path=wallet_file)

    assert wallet_file.exists()
    assert wallet_file.stat().st_mode & 0o777 == 0o600
    assert len(secret) == 32
    assert len(public) == 32
    assert len(address) > 20  # base58-encoded 32 bytes is 43-44 chars typically
    # Solana addresses are base58 — no '0', 'O', 'I', 'l'
    for ch in "0OIl":
        assert ch not in address, f"Invalid base58 char {ch!r} in address {address!r}"


def test_wallet_loads_existing_keypair(tmp_path):
    """Loading the same wallet file twice returns identical keys and address."""
    wallet_file = tmp_path / "wallet.json"

    secret1, public1, address1 = load_or_create_wallet(path=wallet_file)
    secret2, public2, address2 = load_or_create_wallet(path=wallet_file)

    assert secret1 == secret2
    assert public1 == public2
    assert address1 == address2


def test_wallet_load_repairs_insecure_permissions(tmp_path):
    """Existing private key files are tightened to owner-only permissions."""
    wallet_file = tmp_path / "wallet.json"
    load_or_create_wallet(path=wallet_file)
    wallet_file.chmod(0o644)

    load_or_create_wallet(path=wallet_file)

    assert wallet_file.stat().st_mode & 0o777 == 0o600


def test_base58_counts_only_leading_zero_bytes():
    """Zero bytes inside the public key do not become extra base58 leading ones."""
    assert _b58encode(bytes([0, 1, 0])) == "15R"


def test_download_shard_stub_creates_cache(tmp_path):
    """Stub-model shard creates a local cache file without network access."""
    shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path)
    assert shard_dir.exists()
    weights = shard_dir / "weights.json"
    assert weights.exists()
    data = json.loads(weights.read_text())
    assert data["stub"] is True
    assert data["shard_start"] == 0
    assert data["shard_end"] == 31


def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeypatch):
    """Non-stub shards use the HuggingFace snapshot_download path."""
    calls = []

    def fake_snapshot_download(repo_id, cache_dir, local_dir):
        calls.append({"repo_id": repo_id, "cache_dir": cache_dir, "local_dir": local_dir})
        Path(local_dir).mkdir(parents=True, exist_ok=True)
        return local_dir

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        0,
        3,
        cache_dir=tmp_path,
        hf_repo="org/tiny-llama-shards",
        progress=False,
    )

    assert shard_dir == tmp_path / "tiny-llama"
    assert calls == [{
        "repo_id": "org/tiny-llama-shards",
        "cache_dir": str(tmp_path),
        "local_dir": str(shard_dir),
    }]


def test_download_shard_reuses_model_cache_for_narrower_layer_range(
    tmp_path,
    monkeypatch,
):
    """A wider cached shard satisfies a later narrower assignment for the same model."""
    cache_dir = tmp_path / "cache"
    model_dir = cache_dir / "tiny-llama"
    model_dir.mkdir(parents=True)
    (model_dir / "config.json").write_bytes(b"{}")
    (model_dir / "model-00001-of-00002.safetensors").write_bytes(b"a" * 3)
    (model_dir / "model-00002-of-00002.safetensors").write_bytes(b"b" * 5)

    def unexpected_urlopen(*args, **kwargs):
        raise AssertionError("cached files should avoid tracker download")

    def unexpected_snapshot_download(*args, **kwargs):
        raise AssertionError("cached files should avoid HuggingFace download")

    monkeypatch.setattr(urllib.request, "urlopen", unexpected_urlopen)
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=unexpected_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        0,
        1,
        cache_dir=cache_dir,
        hf_repo="org/tiny-llama-shards",
        model_sources=[{
            "type": "tracker",
            "url": "http://tracker/v1/model-files/download?model=tiny-llama",
            "files": ["config.json", "model-00001-of-00002.safetensors"],
            "file_sizes": {
                "config.json": 2,
                "model-00001-of-00002.safetensors": 3,
            },
        }],
        peers=[{"endpoint": "http://peer", "checksum": "unused"}],
        progress=False,
    )

    assert shard_dir == model_dir
    assert (model_dir / "model-00002-of-00002.safetensors").read_bytes() == b"b" * 5


def test_download_shard_prefers_tracker_model_source_over_huggingface(
    tmp_path,
    monkeypatch,
):
    """A working tracker model source is used exclusively — HF is never contacted."""
    contents = {
        "config.json": b"{}",
        "model-00002-of-00004.safetensors": b"tracker",
    }

    class FakeFileResponse:
        def __init__(self, payload: bytes):
            self._payload = io.BytesIO(payload)
            self._length = len(payload)

        def __enter__(self):
            return self

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

        def getheader(self, name: str):
            if name == "Content-Length":
                return str(self._length)
            if name == "Content-Type":
                return "application/octet-stream"
            return None

        def read(self, size: int = -1) -> bytes:
            return self._payload.read(size)

    def fake_urlopen(url, *args, **kwargs):
        query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
        rel = query.get("file", [None])[0]
        assert rel in contents, f"unexpected per-file request: {url}"
        return FakeFileResponse(contents[rel])

    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
    hf_calls = []

    def fake_snapshot_download(**kwargs):
        hf_calls.append(kwargs)
        time.sleep(0.05)
        local_dir = Path(kwargs["local_dir"])
        local_dir.mkdir(parents=True, exist_ok=True)
        (local_dir / "model-00002-of-00004.safetensors").write_text("hf")
        return str(local_dir)

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        2,
        3,
        cache_dir=tmp_path / "cache",
        hf_repo="org/tiny-llama-shards",
        model_sources=[{
            "type": "tracker",
            "url": "http://tracker/v1/model-files/download?model=tiny-llama",
            "files": ["config.json", "model-00002-of-00004.safetensors"],
        }],
        progress=False,
    )

    assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "tracker"
    assert hf_calls == []


def test_download_shard_prefers_tracker_full_model_source_over_huggingface(
    tmp_path,
    monkeypatch,
):
    """A tracker-advertised full snapshot is sufficient on its own — HF is never contacted."""
    contents = {
        "config.json": b"{}",
        "weights-a.safetensors": b"tracker-a",
        "weights-b.safetensors": b"tracker-b",
    }

    class FakeFileResponse:
        def __init__(self, payload: bytes):
            self._payload = io.BytesIO(payload)
            self._length = len(payload)

        def __enter__(self):
            return self

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

        def getheader(self, name: str):
            if name == "Content-Length":
                return str(self._length)
            if name == "Content-Type":
                return "application/octet-stream"
            return None

        def read(self, size: int = -1) -> bytes:
            return self._payload.read(size)

    def fake_urlopen(url, *args, **kwargs):
        query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
        rel = query.get("file", [None])[0]
        assert rel in contents, f"unexpected per-file request: {url}"
        return FakeFileResponse(contents[rel])

    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
    hf_calls = []

    def fake_snapshot_download(**kwargs):
        hf_calls.append(kwargs)
        raise AssertionError("HuggingFace should not be contacted when tracker full_files are available")

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        0,
        3,
        cache_dir=tmp_path / "cache",
        hf_repo="org/tiny-llama-shards",
        model_sources=[{
            "type": "tracker-full",
            "url": "http://tracker/v1/model-files/download?model=tiny-llama&full=1",
            "files": ["config.json", "weights-a.safetensors", "weights-b.safetensors"],
            "full_files": ["config.json", "weights-a.safetensors", "weights-b.safetensors"],
        }],
        progress=False,
    )

    assert (shard_dir / "config.json").read_text() == "{}"
    assert (shard_dir / "weights-a.safetensors").read_text() == "tracker-a"
    assert (shard_dir / "weights-b.safetensors").read_text() == "tracker-b"
    assert hf_calls == []


def test_download_shard_falls_back_to_huggingface_when_tracker_source_fails(
    tmp_path,
    monkeypatch,
):
    """A dead tracker source falls through to HF with allow_patterns from the source files."""

    def failing_urlopen(*args, **kwargs):
        raise ConnectionResetError("tracker went away")

    monkeypatch.setattr(urllib.request, "urlopen", failing_urlopen)
    hf_calls = []

    def fake_snapshot_download(**kwargs):
        hf_calls.append(kwargs)
        local_dir = Path(kwargs["local_dir"])
        local_dir.mkdir(parents=True, exist_ok=True)
        (local_dir / "model-00002-of-00004.safetensors").write_text("hf")
        return str(local_dir)

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        2,
        3,
        cache_dir=tmp_path / "cache",
        hf_repo="org/tiny-llama-shards",
        model_sources=[{
            "type": "tracker",
            "url": "http://tracker/v1/model-files/download?model=tiny-llama",
            "files": ["config.json", "model-00002-of-00004.safetensors"],
        }],
        progress=False,
    )

    assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "hf"
    assert hf_calls[0]["allow_patterns"] == ["config.json", "model-00002-of-00004.safetensors"]


def test_download_shard_logs_huggingface_source(tmp_path, monkeypatch, capsys):
    """Shard download status tells the node operator when HuggingFace was used."""

    def fake_snapshot_download(repo_id, cache_dir, local_dir):
        Path(local_dir).mkdir(parents=True, exist_ok=True)
        (Path(local_dir) / "weights.json").write_text(json.dumps({"repo_id": repo_id}))
        return local_dir

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    download_shard(
        "tiny-llama",
        0,
        3,
        cache_dir=tmp_path,
        hf_repo="org/tiny-llama-shards",
    )

    assert "download source: HuggingFace" in capsys.readouterr().out


def test_download_shard_rejects_peer_checksum_mismatch_before_fallback(
    tmp_path,
    monkeypatch,
):
    """Corrupt peer chunks are not marked complete; HuggingFace remains the fallback."""
    corrupt_dir = tmp_path / "corrupt"
    corrupt_dir.mkdir()
    (corrupt_dir / "weights.json").write_text(json.dumps({"payload": "corrupt"}))
    archive = io.BytesIO()
    write_shard_archive(corrupt_dir, archive)

    class FakePeerResponse:
        def __init__(self, payload: bytes):
            self._payload = io.BytesIO(payload)

        def __enter__(self):
            return self

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

        def read(self, size: int = -1) -> bytes:
            return self._payload.read(size)

    monkeypatch.setattr(
        urllib.request,
        "urlopen",
        lambda *args, **kwargs: FakePeerResponse(archive.getvalue()),
    )
    hf_calls = []

    def fake_snapshot_download(repo_id, cache_dir, local_dir):
        hf_calls.append(repo_id)
        shard_dir = Path(local_dir)
        shard_dir.mkdir(parents=True, exist_ok=True)
        (shard_dir / "weights.json").write_text(json.dumps({"payload": "hf"}))
        return local_dir

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    shard_dir = download_shard(
        "tiny-llama",
        0,
        3,
        cache_dir=tmp_path / "cache",
        hf_repo="org/tiny-llama-shards",
        peers=[{"endpoint": "http://peer", "checksum": "not-the-corrupt-checksum"}],
        progress=False,
    )

    assert hf_calls == ["org/tiny-llama-shards"]
    assert json.loads((shard_dir / "weights.json").read_text()) == {"payload": "hf"}


def test_download_shard_stub_idempotent(tmp_path):
    """Calling download_shard twice does not error — file already exists."""
    download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False)
    shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False)
    assert shard_dir.exists()


def test_startup_formats_probationary_jobs_remaining():
    """Startup status tells a node how many free jobs remain before earning."""
    contracts = LocalSolanaContracts(probationary_job_count=50)
    for _ in range(12):
        contracts.registry.record_completed_job("node-wallet-a")

    line = _probationary_status_line(contracts, "node-wallet-a")

    assert line == "Probationary period: 38 jobs remaining before earning"


# ---------------------------------------------------------------------------
# Tracker assign endpoint
# ---------------------------------------------------------------------------


def _get_json(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=5) as r:
        return json.loads(r.read())


def test_tracker_assign_returns_shard_for_empty_registry():
    """Tracker assigns the full layer range when no nodes are registered."""
    tracker = TrackerServer()
    port = tracker.start()
    try:
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0"
        )
        assert resp["shard_start"] == 0
        assert resp["shard_end"] == 15
        assert resp["model"] == "stub-model"
        assert resp["model_layers_end"] == 31
    finally:
        tracker.stop()


def test_tracker_assign_fills_gap():
    """Tracker assigns the first uncovered layer range when a node is already registered."""
    import json as _json
    import urllib.request as _ur

    tracker = TrackerServer()
    port = tracker.start()
    try:
        # Register a node covering layers 0-15
        data = _json.dumps({
            "endpoint": "http://127.0.0.1:9100",
            "shard_start": 0,
            "shard_end": 15,
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req = _ur.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with _ur.urlopen(req) as r:
            r.read()

        # Assignment should fill the gap: layers 16-31
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0"
        )
        assert resp["shard_start"] == 16
        assert resp["shard_end"] == 31
    finally:
        tracker.stop()


def test_tracker_assign_returns_huggingface_repo_when_configured():
    """Tracker includes the HuggingFace repo identifier in shard assignments."""
    tracker = TrackerServer(model_presets={
        "tiny-llama": {"layers_start": 0, "layers_end": 7, "hf_repo": "org/tiny-llama-shards"}
    })
    port = tracker.start()
    try:
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cuda&vram_mb=24576"
        )
        assert resp["model"] == "tiny-llama"
        assert resp["hf_repo"] == "org/tiny-llama-shards"
        assert resp["shard_start"] == 0
        assert resp["shard_end"] == 7
    finally:
        tracker.stop()


def test_tracker_assign_advertises_local_model_source_and_serves_subset(tmp_path):
    """Tracker with models_dir advertises and serves only files needed for the shard."""
    snapshot = tmp_path / "models" / "models--org--tiny-llama-shards" / "snapshots" / "abc"
    nested = snapshot / "nested"
    nested.mkdir(parents=True)
    (snapshot / "config.json").write_text(json.dumps({"num_hidden_layers": 4}))
    (snapshot / "tokenizer.json").write_text("{}")
    (snapshot / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
            "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
            "model.layers.1.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
            "model.layers.2.self_attn.q_proj.weight": "nested/model-00002-of-00003.safetensors",
            "model.layers.3.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
            "lm_head.weight": "model-00003-of-00003.safetensors",
        },
    }))
    for rel in [
        "model-00001-of-00003.safetensors",
        "model-00002-of-00003.safetensors",
        "nested/model-00002-of-00003.safetensors",
        "model-00003-of-00003.safetensors",
    ]:
        (snapshot / rel).write_text(rel)

    tracker = TrackerServer(
        model_presets={
            "tiny-llama": {
                "layers_start": 0,
                "layers_end": 3,
                "hf_repo": "org/tiny-llama-shards",
                "bytes_per_layer": {"bfloat16": 1024 * 1024},
            },
        },
        models_dir=tmp_path / "models",
    )
    port = tracker.start()
    try:
        data = json.dumps({
            "endpoint": "http://127.0.0.1:9100",
            "model": "tiny-llama",
            "shard_start": 0,
            "shard_end": 0,
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            r.read()
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cpu&ram_mb=3"
        )
        assert resp["shard_start"] == 1
        assert resp["shard_end"] == 2
        assert resp["model_sources"]
        source = resp["model_sources"][0]
        assert source["files"] == [
            "config.json",
            "model-00002-of-00003.safetensors",
            "model.safetensors.index.json",
            "nested/model-00002-of-00003.safetensors",
            "tokenizer.json",
        ]
        with urllib.request.urlopen(source["url"], timeout=5) as response:
            payload = io.BytesIO(response.read())
        with tarfile.open(fileobj=payload, mode="r") as tf:
            names = sorted(tf.getnames())
        assert names == source["files"]
    finally:
        tracker.stop()


def test_tracker_assign_lists_peers_for_same_model_shard():
    """A registered node with a completed shard is returned as a same-shard peer."""
    import json as _json
    import urllib.request as _ur

    tracker = TrackerServer(model_presets={
        "tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
    })
    port = tracker.start()
    try:
        data = _json.dumps({
            "endpoint": "http://127.0.0.1:9100",
            "model": "tiny-llama",
            "shard_start": 0,
            "shard_end": 15,
            "shard_checksum": "abc123",
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req = _ur.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with _ur.urlopen(req) as r:
            r.read()

        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cpu&vram_mb=0"
        )

        assert resp["shard_start"] == 0
        assert resp["shard_end"] == 15
        assert resp["hf_repo"] == "org/tiny-llama-shards"
        assert resp["peers"] == [{
            "endpoint": "http://127.0.0.1:9100",
            "checksum": "abc123",
        }]
    finally:
        tracker.stop()


def test_infer_relay_url_from_public_https_tracker():
    assert _infer_relay_url_from_tracker("https://ai.neuron.d-popov.com") == (
        "wss://ai.neuron.d-popov.com/ws"
    )
    assert _infer_relay_url_from_tracker("https://ai.neuron.d-popov.com/v1/network/map") == (
        "wss://ai.neuron.d-popov.com/ws"
    )
    assert _infer_relay_url_from_tracker("http://192.168.0.179:8081") is None
    assert _infer_relay_url_from_tracker("http://127.0.0.1:8081") is None


def test_public_https_tracker_infers_relay_when_network_map_omits_relay_url(
    tmp_path,
    monkeypatch,
    capsys,
):
    """Nodes bootstrap relay from the tracker origin when map relay_url is null."""
    import meshnet_node.startup as startup_mod

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            self.backend = FakeBackend()
            self.port = None
            self.chat_completion_count = 0
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8001
            return self.port

        def stop(self):
            pass

    class FakeRelayHttpBridge:
        def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
            self.relay_url = relay_url
            self.peer_id = peer_id

        @property
        def relay_addr(self):
            return f"{self.relay_url.replace('/ws', '')}/rpc/{self.peer_id}"

        def start(self):
            return types.SimpleNamespace(peer_id=self.peer_id, relay_addr=self.relay_addr)

        def wait_connected(self, timeout=5.0):
            return True

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
    monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)
    monkeypatch.setattr(
        startup_mod,
        "_get_json",
        lambda url, timeout=10.0: {"relay_url": None, "nodes": []},
    )

    tracker_url = "https://ai.neuron.d-popov.com"
    node = run_startup(
        tracker_url=tracker_url,
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        advertise_host="172.29.104.23",
        wallet_path=tmp_path / "wallet.json",
    )
    try:
        pass
    finally:
        node.stop()

    output = capsys.readouterr().out
    assert "Relay advertised by tracker" in output
    assert "wss://ai.neuron.d-popov.com/ws" in output
    assert "Cross-host pipeline hops WILL time out" not in output


def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, capsys):
    """Real-model startup summary prints the shard range plus total model layers."""
    import meshnet_node.startup as startup_mod
    captured_registration = {}

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            self.kwargs = kwargs
            self.backend = FakeBackend()
            self.port = None

        def start(self):
            self.port = 8001
            return self.port

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
    monkeypatch.setattr(
        startup_mod,
        "_post_json",
        lambda _url, _payload, timeout=10.0: (
            captured_registration.update(_payload) or {"node_id": "node-test-123"}
        ),
    )

    node = run_startup(
        tracker_url="http://127.0.0.1:8080",
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        shard_start=0,
        shard_end=23,
        vram_mb_override=6144,
        max_loaded_shards=2,
        wallet_path=tmp_path / "wallet.json",
    )

    assert node.backend.total_layers == 24
    assert node.tracker_node_id == "node-test-123"
    assert captured_registration["vram_bytes"] == 6144 * 1024 * 1024
    assert captured_registration["max_loaded_shards"] == 2
    output = capsys.readouterr().out
    assert "Shard:        layers 0–23; 24 of 24" in output
    assert "Node ID:      node-test-123" in output


def test_real_model_startup_autodetects_cpu_memory_budget_and_logs_shard_budget(
    tmp_path,
    monkeypatch,
    capsys,
):
    """Without --memory, startup reports RAM-backed capacity to the tracker and operator."""
    import meshnet_node.startup as startup_mod
    captured_registration = {}

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            self.kwargs = kwargs
            self.backend = FakeBackend()
            self.port = None
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8001
            return self.port

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
    monkeypatch.setattr(
        startup_mod,
        "_post_json",
        lambda _url, _payload, timeout=10.0: (
            captured_registration.update(_payload) or {"node_id": "node-auto-mem"}
        ),
    )

    node = run_startup(
        tracker_url="http://127.0.0.1:8080",
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        shard_start=0,
        shard_end=23,
        wallet_path=tmp_path / "wallet.json",
    )
    try:
        pass
    finally:
        node.stop()

    assert captured_registration["vram_bytes"] == 0
    assert captured_registration["ram_bytes"] == 16384 * 1024 * 1024
    assert captured_registration["max_loaded_shards"] == 1
    output = capsys.readouterr().out
    assert "Memory budget: 16.0 GB RAM" in output
    assert "Shard budget: up to 24/24 layers at bfloat16" in output
    assert "GB remaining after full load" in output
    assert "Node ID:      node-auto-mem" in output


def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_only(
    tmp_path,
    monkeypatch,
    capsys,
):
    """A node only needs the public tracker URL to discover relay metadata and register."""
    import meshnet_node.startup as startup_mod

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            self.kwargs = kwargs
            self.backend = FakeBackend()
            self.port = None
            self.chat_completion_count = 0
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8001
            return self.port

        def stop(self):
            pass

    class FakeRelayHttpBridge:
        def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
            self.relay_url = relay_url
            self.peer_id = peer_id
            self.local_base_url = local_base_url
            self.advertised_addr = advertised_addr

        @property
        def relay_addr(self):
            return f"ws://public-relay.example/rpc/{self.peer_id}"

        def start(self):
            return types.SimpleNamespace(
                peer_id=self.peer_id,
                relay_addr=self.relay_addr,
            )

        def wait_connected(self, timeout=5.0):
            return True

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
    monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)

    tracker = TrackerServer(relay_url="ws://public-relay.example/ws")
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        node = run_startup(
            tracker_url=tracker_url,
            model_id="Qwen/Qwen2.5-0.5B-Instruct",
            wallet_path=tmp_path / "wallet.json",
        )
        try:
            network_map = _get_json(f"{tracker_url}/v1/network/map")
        finally:
            node.stop()
    finally:
        tracker.stop()

    assert network_map["relay_url"] == "ws://public-relay.example/ws"
    assert len(network_map["nodes"]) == 1
    registered = network_map["nodes"][0]
    assert registered["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
    assert registered["endpoint"].startswith("http://")
    assert registered["endpoint"].endswith(":8001")
    assert registered["relay_addr"].startswith("ws://public-relay.example/rpc/")
    assert registered["peer_id"]
    output = capsys.readouterr().out
    assert "Relay advertised by tracker" in output
    assert "Cross-host pipeline hops WILL time out" not in output


def test_public_tracker_relay_suppresses_virtual_ip_warning(
    tmp_path,
    monkeypatch,
    capsys,
):
    """A WSL/Docker endpoint is acceptable when the tracker advertises relay RPC."""
    import meshnet_node.startup as startup_mod

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            self.backend = FakeBackend()
            self.port = None
            self.chat_completion_count = 0
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8001
            return self.port

        def stop(self):
            pass

    class FakeRelayHttpBridge:
        def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
            self.relay_url = relay_url
            self.peer_id = peer_id

        @property
        def relay_addr(self):
            return f"ws://public-relay.example/rpc/{self.peer_id}"

        def start(self):
            return types.SimpleNamespace(peer_id=self.peer_id, relay_addr=self.relay_addr)

        def wait_connected(self, timeout=5.0):
            return True

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
    monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)

    tracker = TrackerServer(relay_url="ws://public-relay.example/ws")
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        node = run_startup(
            tracker_url=tracker_url,
            model_id="Qwen/Qwen2.5-0.5B-Instruct",
            advertise_host="172.29.104.23",
            wallet_path=tmp_path / "wallet.json",
        )
        try:
            network_map = _get_json(f"{tracker_url}/v1/network/map")
        finally:
            node.stop()
    finally:
        tracker.stop()

    assert network_map["nodes"][0]["endpoint"] == "http://172.29.104.23:8001"
    assert network_map["nodes"][0]["relay_addr"].startswith("ws://public-relay.example/rpc/")
    output = capsys.readouterr().out
    assert "Relay advertised by tracker" in output
    assert "Cross-host pipeline hops WILL time out" not in output


def test_later_node_auto_joins_existing_public_hf_model_with_only_tracker_url(
    tmp_path,
    monkeypatch,
):
    """After a model exists, a node can join by knowing only the public tracker URL."""
    import meshnet_node.startup as startup_mod

    captured = {}

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            captured.update(kwargs)
            self.backend = FakeBackend()
            self.port = None
            self.chat_completion_count = 0
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8002
            return self.port

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)

    tracker = TrackerServer()
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        data = json.dumps({
            "endpoint": "http://203.0.113.20:8001",
            "model": "Qwen2.5-0.5B-Instruct",
            "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
            "num_layers": 24,
            "shard_start": 0,
            "shard_end": 11,
            "tracker_mode": True,
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req = urllib.request.Request(
            f"{tracker_url}/v1/nodes/register",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req) as resp:
            resp.read()

        node = run_startup(
            tracker_url=tracker_url,
            advertise_host="203.0.113.21",
            wallet_path=tmp_path / "wallet.json",
        )
        try:
            route_resp = _get_json(
                f"{tracker_url}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
            )
        finally:
            node.stop()
    finally:
        tracker.stop()

    assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
    assert captured["shard_start"] == 12
    assert captured["shard_end"] == 23
    assert route_resp["route"] == ["http://203.0.113.20:8001", "http://203.0.113.21:8002"]


def test_later_node_auto_joins_redundant_copy_when_model_is_fully_covered(
    tmp_path,
    monkeypatch,
):
    """Model-less joins should load the served HF model even when gap_found=false."""
    import meshnet_node.startup as startup_mod

    captured = {}

    class FakeBackend:
        total_layers = 24

    class FakeTorchNodeServer:
        def __init__(self, **kwargs):
            captured.update(kwargs)
            self.backend = FakeBackend()
            self.port = None
            self.chat_completion_count = 0
            self.total_requests = 0
            self.failed_requests = 0
            self.queue_depth = 0

        def start(self):
            self.port = 8003
            return self.port

        def stop(self):
            pass

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)

    tracker = TrackerServer()
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        for endpoint, shard_start, shard_end in (
            ("http://203.0.113.30:8001", 0, 11),
            ("http://203.0.113.31:8001", 12, 23),
        ):
            data = json.dumps({
                "endpoint": endpoint,
                "model": "Qwen2.5-0.5B-Instruct",
                "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
                "num_layers": 24,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "tracker_mode": shard_start == 0,
                "hardware_profile": {},
                "score": 1.0,
            }).encode()
            req = urllib.request.Request(
                f"{tracker_url}/v1/nodes/register",
                data=data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req) as resp:
                resp.read()

        node = run_startup(
            tracker_url=tracker_url,
            advertise_host="203.0.113.32",
            wallet_path=tmp_path / "wallet.json",
        )
        try:
            assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
            assert captured["shard_start"] == 0
        finally:
            node.stop()
    finally:
        tracker.stop()


# ---------------------------------------------------------------------------
# Full startup integration test
# ---------------------------------------------------------------------------


def test_full_startup_sequence(tmp_path):
    """Full startup: hardware → wallet → assign → download → start → register."""
    tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    wallet_path = tmp_path / "wallet.json"
    cache_dir = tmp_path / "shards"

    try:
        node = run_startup(
            tracker_url=tracker_url,
            model="stub-model",
            wallet_path=wallet_path,
            cache_dir=cache_dir,
        )
        try:
            # Wallet was created on disk
            assert wallet_path.exists()
            wallet_data = json.loads(wallet_path.read_text())
            assert len(wallet_data) == 64

            # Shard was cached
            assert any(cache_dir.rglob("weights.json"))

            # Node appears in tracker registry (route resolves successfully)
            route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model")
            assert len(route_resp["route"]) >= 1
            assert node.port is not None
            assert any(str(node.port) in ep for ep in route_resp["route"])

            # Node accepts an inference request
            payload = json.dumps({
                "messages": [{"role": "user", "content": "hello"}],
            }).encode()
            req = urllib.request.Request(
                f"http://127.0.0.1:{node.port}/v1/infer",
                data=payload,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=5) as resp:
                assert resp.status == 200
                body = json.loads(resp.read())
            # Last-shard node returns text; first-shard returns activations
            assert "text" in body or "activations" in body
        finally:
            node.stop()
    finally:
        tracker.stop()


def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch):
    import meshnet_node.startup as startup_mod

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    heartbeat_calls = []
    monkeypatch.setattr(
        startup_mod,
        "_start_heartbeat",
        lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)),
    )

    tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        node = run_startup(
            tracker_url=tracker_url,
            model="stub-model",
            wallet_path=tmp_path / "wallet.json",
            cache_dir=tmp_path / "shards",
        )
        try:
            assert len(heartbeat_calls) == 1
            args, kwargs = heartbeat_calls[0]
            assert args[0] == tracker_url
            assert args[2]["model"] == "stub-model"
            assert kwargs["node_ref"] is node
        finally:
            node.stop()
    finally:
        tracker.stop()


def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch):
    """Explicit --shard-start/--shard-end override tracker auto-assignment."""
    import meshnet_node.startup as startup_mod

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024},
    )
    heartbeat_calls = []
    monkeypatch.setattr(
        startup_mod,
        "_start_heartbeat",
        lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)),
    )

    tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        node = run_startup(
            tracker_url=tracker_url,
            model="stub-model",
            shard_start=0,
            shard_end=5,
            wallet_path=tmp_path / "wallet.json",
            cache_dir=tmp_path / "shards",
        )
        try:
            assert len(heartbeat_calls) == 1
            args, kwargs = heartbeat_calls[0]
            reg_payload = args[2]
            assert reg_payload["shard_start"] == 0
            assert reg_payload["shard_end"] == 5
            assert reg_payload["managed_assignment"] is False
        finally:
            node.stop()
    finally:
        tracker.stop()


def test_real_model_startup_registers_downloaded_inventory_without_checksum(
    tmp_path,
    monkeypatch,
    capsys,
):
    """Real model folders are reported as inventory without hashing their contents."""
    import meshnet_node.startup as startup_mod

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    monkeypatch.setattr(
        startup_mod,
        "compute_shard_checksum",
        lambda _path: (_ for _ in ()).throw(AssertionError("real model startup must not hash model files")),
    )
    hf_calls = []

    def fake_snapshot_download(repo_id, cache_dir, local_dir):
        hf_calls.append({"repo_id": repo_id, "local_dir": local_dir})
        shard_dir = Path(local_dir)
        shard_dir.mkdir(parents=True, exist_ok=True)
        (shard_dir / "weights.json").write_text(json.dumps({
            "repo_id": repo_id,
            "payload": "node-a-hf-stub",
        }))
        return local_dir

    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
    )

    tracker = TrackerServer(model_presets={
        "tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
    })
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"
    try:
        node = run_startup(
            tracker_url=tracker_url,
            model="tiny-llama",
            wallet_path=tmp_path / "wallet.json",
            cache_dir=tmp_path / "node-shards",
        )
        try:
            assert len(hf_calls) == 1
            assert (tmp_path / "node-shards" / "tiny-llama" / "weights.json").exists()
            output = capsys.readouterr().out
            assert "Cached at:" in output
            network_map = _get_json(f"{tracker_url}/v1/network/map")
            registered = network_map["nodes"][0]
            assert registered["downloaded_models"] == [{
                "model": "tiny-llama",
                "shard_start": 0,
                "shard_end": 15,
                "path": str(tmp_path / "node-shards" / "tiny-llama"),
                "file_count": 1,
                "total_bytes": (tmp_path / "node-shards" / "tiny-llama" / "weights.json").stat().st_size,
                "hf_repo": "org/tiny-llama-shards",
            }]
        finally:
            node.stop()
    finally:
        tracker.stop()


def test_downloaded_model_inventory_reports_local_model_percentage(tmp_path):
    import meshnet_node.startup as startup_mod

    model_dir = tmp_path / "models" / "tiny-llama"
    model_dir.mkdir(parents=True)
    (model_dir / "config.json").write_bytes(b"{}")
    (model_dir / "weights-a.safetensors").write_bytes(b"a" * 3)

    inventory = startup_mod._downloaded_model_inventory(
        "tiny-llama",
        0,
        1,
        model_dir,
        hf_repo="org/tiny-llama",
        model_sources=[{
            "full_files": ["config.json", "weights-a.safetensors", "weights-b.safetensors"],
            "file_sizes": {
                "config.json": 2,
                "weights-a.safetensors": 3,
                "weights-b.safetensors": 5,
            },
        }],
    )

    assert inventory == [{
        "model": "tiny-llama",
        "shard_start": 0,
        "shard_end": 1,
        "path": str(model_dir),
        "file_count": 2,
        "total_bytes": 5,
        "hf_repo": "org/tiny-llama",
        "expected_file_count": 3,
        "local_expected_file_count": 2,
        "expected_bytes": 10,
        "local_expected_bytes": 5,
        "local_model_percentage": 50.0,
    }]


def test_network_assign_gap_found_field():
    """network/assign sets gap_found=True when a real gap exists, False when fully covered."""
    import json as _json
    import urllib.request as _ur

    tracker = TrackerServer()
    port = tracker.start()
    try:
        # Register a node covering only layers 0-11 of a 24-layer model.
        data = _json.dumps({
            "endpoint": "http://127.0.0.1:9200",
            "model": "Qwen2.5-0.5B-Instruct",
            "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
            "num_layers": 24,
            "shard_start": 0,
            "shard_end": 11,
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req = _ur.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with _ur.urlopen(req) as r:
            r.read()

        # A new node should be told there is a gap (layers 12-23).
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
            "&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
        )
        assert resp["gap_found"] is True
        assert resp["shard_start"] == 12, f"expected gap at 12, got {resp['shard_start']}"
        assert resp["shard_end"] == 23

        # Register the second node covering the gap.
        data2 = _json.dumps({
            "endpoint": "http://127.0.0.1:9201",
            "model": "Qwen2.5-0.5B-Instruct",
            "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
            "num_layers": 24,
            "shard_start": 12,
            "shard_end": 23,
            "hardware_profile": {},
            "score": 1.0,
        }).encode()
        req2 = _ur.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=data2,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with _ur.urlopen(req2) as r:
            r.read()

        # Now fully covered — gap_found should be False.
        resp2 = _get_json(
            f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
            "&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
        )
        assert resp2["gap_found"] is False
    finally:
        tracker.stop()


def test_route_finds_hf_model_across_two_nodes():
    """Tracker /v1/route returns ordered route for HF model even without a preset."""
    import json as _json
    import urllib.request as _ur

    tracker = TrackerServer()
    port = tracker.start()
    try:
        def register(endpoint, shard_start, shard_end):
            data = _json.dumps({
                "endpoint": endpoint,
                "model": "Qwen2.5-0.5B-Instruct",
                "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
                "num_layers": 24,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "hardware_profile": {},
                "score": 1.0,
            }).encode()
            req = _ur.Request(
                f"http://127.0.0.1:{port}/v1/nodes/register",
                data=data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with _ur.urlopen(req) as r:
                r.read()

        register("http://127.0.0.1:9300", 0, 11)
        register("http://127.0.0.1:9301", 12, 23)

        # Route by hf_repo (full identifier).
        resp = _get_json(
            f"http://127.0.0.1:{port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
        )
        assert resp["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]

        # Route also works by short model name.
        resp2 = _get_json(
            f"http://127.0.0.1:{port}/v1/route?model=Qwen2.5-0.5B-Instruct"
        )
        assert resp2["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]
    finally:
        tracker.stop()


def test_register_deduplicates_same_endpoint():
    """Re-registering the same endpoint replaces the old entry, not duplicates it."""
    import json as _json
    import urllib.request as _ur

    tracker = TrackerServer()
    port = tracker.start()
    try:
        def register(shard_start, shard_end):
            data = _json.dumps({
                "endpoint": "http://127.0.0.1:9400",
                "model": "Qwen2.5-0.5B-Instruct",
                "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
                "num_layers": 24,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "hardware_profile": {},
                "score": 1.0,
            }).encode()
            req = _ur.Request(
                f"http://127.0.0.1:{port}/v1/nodes/register",
                data=data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with _ur.urlopen(req) as r:
                return _json.loads(r.read())

        register(0, 23)  # initial full-model registration
        register(12, 23)  # re-register with corrected shard range

        # After re-register, tracker should see only one node at 12-23 for this endpoint.
        # If both were still registered, the gap scan would find no gap (0-23 still covers).
        # With dedup, the old 0-23 is gone and a real gap 0-11 exists.
        assign_resp = _get_json(
            f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
            "&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
        )
        assert assign_resp["gap_found"] is True
        assert assign_resp["shard_start"] == 0, "old 0-23 entry should have been replaced"
    finally:
        tracker.stop()


def test_startup_cpu_fallback(tmp_path, monkeypatch):
    """Node starts with CPU warning when no GPU is detected."""
    import meshnet_node.startup as startup_mod

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )

    tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    try:
        node = run_startup(
            tracker_url=tracker_url,
            model="stub-model",
            wallet_path=tmp_path / "wallet.json",
            cache_dir=tmp_path / "shards",
        )
        try:
            # Node is running even on CPU
            assert node.port is not None
            route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model")
            assert len(route_resp["route"]) >= 1
        finally:
            node.stop()
    finally:
        tracker.stop()


# --------------------------------------------------- layer detection (US: composite configs)


def test_detect_num_layers_prefers_flattened_local_model_config(tmp_path, monkeypatch):
    import meshnet_node.startup as startup_mod

    model_dir = tmp_path / "Qwen3.6-35B-A3B"
    model_dir.mkdir()
    (model_dir / "config.json").write_text("{}")
    calls = []

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            calls.append({"model_id": model_id, "cache_dir": cache_dir})
            return types.SimpleNamespace(num_hidden_layers=37)

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        types.SimpleNamespace(AutoConfig=AutoConfigStub),
    )

    assert startup_mod._detect_num_layers("unsloth/Qwen3.6-35B-A3B", cache_dir=tmp_path) == 37
    assert calls == [{"model_id": str(model_dir), "cache_dir": None}]


def test_layers_from_config_top_level():
    from meshnet_node.model_catalog import layers_from_config

    cfg = types.SimpleNamespace(num_hidden_layers=24)
    assert layers_from_config(cfg) == 24


def test_layers_from_config_nested_text_config():
    """VLM/MoE composites (e.g. Qwen3.5-MoE) keep the layer count in text_config."""
    from meshnet_node.model_catalog import layers_from_config

    cfg = types.SimpleNamespace(text_config=types.SimpleNamespace(num_hidden_layers=40))
    assert layers_from_config(cfg) == 40


def test_layers_from_config_get_text_config_and_variants():
    from meshnet_node.model_catalog import layers_from_config

    inner = types.SimpleNamespace(n_layer=32)
    cfg = types.SimpleNamespace(get_text_config=lambda: inner)
    assert layers_from_config(cfg) == 32
    assert layers_from_config(types.SimpleNamespace()) is None


def test_download_dir_env_override(tmp_path, monkeypatch):
    import importlib

    from meshnet_node import config as config_mod

    monkeypatch.chdir(tmp_path)
    monkeypatch.setenv("MESHNET_DOWNLOAD_DIR", "/tmp/llm-store")
    importlib.reload(config_mod)
    assert config_mod.DEFAULTS["download_dir"] == "/tmp/llm-store"
    monkeypatch.delenv("MESHNET_DOWNLOAD_DIR")
    importlib.reload(config_mod)
    assert config_mod.DEFAULTS["download_dir"].endswith("models")


def test_cli_loads_local_env_before_config_defaults(tmp_path, monkeypatch):
    import importlib

    from meshnet_node import cli as cli_mod
    from meshnet_node import config as config_mod

    monkeypatch.delenv("MESHNET_DOWNLOAD_DIR", raising=False)
    monkeypatch.chdir(tmp_path)
    (tmp_path / ".env").write_text(
        "MESHNET_DOWNLOAD_DIR=/run/media/popov/DATA/llm/safetensor/models\n"
        "HF_TOKEN=hf_test_token\n"
    )

    cli_mod._load_env_defaults()
    importlib.reload(config_mod)

    assert config_mod.DEFAULTS["download_dir"] == "/run/media/popov/DATA/llm/safetensor/models"
    assert os.environ["HF_TOKEN"] == "hf_test_token"


def test_default_quantization_is_auto(monkeypatch):
    import importlib

    from meshnet_node import config as config_mod
    from meshnet_node.model_backend import validate_quantization

    monkeypatch.delenv("MESHNET_DOWNLOAD_DIR", raising=False)
    importlib.reload(config_mod)

    assert config_mod.DEFAULTS["quantization"] == "auto"
    assert validate_quantization("auto") == "auto"


def test_auto_quantization_uses_native_model_dtype_for_unquantized_config():
    from meshnet_node.model_backend import _model_load_plan

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            assert model_id == "repo/model"
            assert cache_dir is None
            return types.SimpleNamespace(
                text_config=types.SimpleNamespace(dtype="torch.bfloat16"),
            )

    torch_stub = types.SimpleNamespace(bfloat16="bf16", float16="fp16")

    quant_config, dtype, uses_quantized_weights = _model_load_plan(
        AutoConfigStub,
        "repo/model",
        "auto",
        torch_stub,
    )

    assert quant_config is None
    assert dtype == "bf16"
    assert uses_quantized_weights is False


def test_auto_quantization_preserves_native_quantized_config():
    from meshnet_node.model_backend import _model_load_plan

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            return types.SimpleNamespace(
                quantization_config={"quant_method": "gptq"},
                torch_dtype="float16",
            )

    torch_stub = types.SimpleNamespace(bfloat16="bf16", float16="fp16")

    quant_config, dtype, uses_quantized_weights = _model_load_plan(
        AutoConfigStub,
        "repo/model",
        "auto",
        torch_stub,
    )

    assert quant_config is None
    assert dtype == "fp16"
    assert uses_quantized_weights is True