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

import json
import io
import sys
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
from meshnet_node.startup import _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)
    if hw["device"] == "cpu":
        assert hw["gpu_name"] is None
        assert hw["vram_mb"] == 0
    else:
        assert isinstance(hw["gpu_name"], str) and hw["gpu_name"]
        assert hw["vram_mb"] > 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" / "layers_0-3"
    assert calls == [{
        "repo_id": "org/tiny-llama-shards",
        "cache_dir": str(tmp_path),
        "local_dir": str(shard_dir),
    }]


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_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_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

    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)

    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",
    )

    assert node.backend.total_layers == 24
    output = capsys.readouterr().out
    assert "Shard:        layers 0–23; 24 of 24" in output


# ---------------------------------------------------------------------------
# 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_second_node_downloads_same_shard_from_peer_without_huggingface(
    tmp_path,
    monkeypatch,
    capsys,
):
    """Node A downloads from HF stub; node B downloads same assignment from node A."""
    import meshnet_node.startup as startup_mod

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
    )
    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}"
    nodes = []
    try:
        node_a = run_startup(
            tracker_url=tracker_url,
            model="tiny-llama",
            wallet_path=tmp_path / "wallet-a.json",
            cache_dir=tmp_path / "node-a-shards",
        )
        nodes.append(node_a)
        assert len(hf_calls) == 1

        node_b = run_startup(
            tracker_url=tracker_url,
            model="tiny-llama",
            wallet_path=tmp_path / "wallet-b.json",
            cache_dir=tmp_path / "node-b-shards",
        )
        nodes.append(node_b)

        assert len(hf_calls) == 1
        assert (tmp_path / "node-b-shards" / "tiny-llama" / "layers_0-15" / "weights.json").exists()
        output = capsys.readouterr().out
        assert "download source: HuggingFace" in output
        assert "download source: peer" in output
    finally:
        for node in reversed(nodes):
            node.stop()
        tracker.stop()


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_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()