"""Full node startup sequence — self-configuring, non-interactive."""

from __future__ import annotations

import json
import socket
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware, benchmark_throughput
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer
from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet


_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024


def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
    """Return the capacity budget in MB and whether it came from VRAM or RAM."""
    if device == "cuda" and vram_mb > 0:
        if shared_vram_mb > 0:
            return vram_mb + shared_vram_mb, "VRAM + shared RAM"
        return vram_mb, "VRAM"
    return max(0, ram_mb), "RAM"


def _hardware_label(device: str, gpu_name: str | None = None) -> str:
    if device == "cuda":
        return "CUDA"
    if gpu_name:
        return "CPU (CUDA inactive)"
    return "CPU"


def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int:
    if total_layers is None or total_layers <= 0 or memory_mb <= 0:
        return 0
    budget_bytes = memory_mb * 1024 * 1024
    return min(total_layers, int((budget_bytes * 0.8) // _DEFAULT_BYTES_PER_LAYER))


def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str:
    memory_gb = memory_mb / 1024
    gb_str = f"{memory_gb:.1f} GB"
    if total_layers is None or total_layers <= 0:
        return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
    max_layers = _max_assignable_layers(memory_mb, total_layers)
    # Remaining capacity after one full model load (rough estimate)
    shard_bytes = max_layers * _DEFAULT_BYTES_PER_LAYER
    remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
    remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
    return (
        f"Memory budget: {gb_str} {memory_source}; "
        f"Shard budget: up to {max_layers}/{total_layers} layers at {quantization}"
        f"{remaining_str}"
    )


def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


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


def _infer_relay_url_from_tracker(tracker_url: str) -> str | None:
    """Infer relay WebSocket URL from a public HTTPS tracker origin.

    Public deployments colocate relay at /ws on the same host as the tracker API
    (see QUICKSTART nginx layout).  Local LAN trackers use a separate relay port
    and must advertise relay_url explicitly via /v1/network/map.
    """
    parsed = urllib.parse.urlparse(tracker_url)
    if parsed.scheme != "https":
        return None
    host = parsed.hostname
    if not host or host in ("127.0.0.1", "localhost"):
        return None
    return f"wss://{parsed.netloc}/ws"


def _discover_relay_url(tracker_url: str) -> str | None:
    relay_url: str | None = None
    try:
        network_map = _get_json(f"{tracker_url}/v1/network/map", timeout=5.0)
        raw = network_map.get("relay_url")
        if isinstance(raw, str) and raw:
            relay_url = raw
    except Exception:
        pass
    return relay_url or _infer_relay_url_from_tracker(tracker_url)


def _start_relay_bridge_if_available(
    tracker_url: str,
    wallet_address: str,
    local_base_url: str,
    advertised_endpoint: str,
    relay_url: str | None = None,
) -> tuple[RelayHttpBridge | None, dict]:
    relay_url = relay_url or _discover_relay_url(tracker_url)
    if not relay_url:
        return None, {}
    peer_id = peer_id_from_wallet(wallet_address)
    bridge = RelayHttpBridge(
        relay_url=relay_url,
        peer_id=peer_id,
        local_base_url=local_base_url,
        advertised_addr=advertised_endpoint,
    )
    info = bridge.start()
    if bridge.wait_connected(timeout=5.0):
        print(f"  Relay connected — {info.relay_addr}", flush=True)
    else:
        print(f"  Relay configured but not connected yet — {info.relay_addr}", flush=True)
    return bridge, {
        "relay_addr": info.relay_addr,
        "peer_id": info.peer_id,
    }


def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHttpBridge | None) -> None:
    setattr(node, "relay_bridge", bridge)
    if bridge is None:
        return
    original_stop = node.stop

    def _stop_with_bridge() -> None:
        try:
            bridge.stop()
        finally:
            original_stop()

    node.stop = _stop_with_bridge  # type: ignore[method-assign]


def _start_heartbeat(
    tracker_url: str,
    node_id: str,
    register_payload: dict,
    interval: float = 20.0,
    node_ref: Any | None = None,
    start_time: float | None = None,
) -> threading.Thread:
    """Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.

    Heartbeat body carries cumulative stats (total_requests, failed_requests,
    queue_depth, uptime_seconds, status). Stats are buffered locally during
    outage and flushed on next successful heartbeat.

    Heartbeat response may include new_assignment: {model, shard_start, shard_end}
    which is logged for now (hot-reload implemented in US-026).
    """
    _start_time = start_time or time.monotonic()

    def _get_stats() -> dict:
        uptime = time.monotonic() - _start_time
        stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
        if node_ref is not None:
            stats["total_requests"] = getattr(
                node_ref,
                "total_requests",
                getattr(node_ref, "chat_completion_count", 0),
            )
            stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
            stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
        return stats

    def _reregister() -> bool:
        nonlocal node_id
        try:
            resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
            node_id = resp.get("node_id", node_id)
            return True
        except Exception:
            return False

    def _apply_directives(directives: list[dict]) -> None:
        if not directives:
            return
        if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
            print(f"  [node] tracker directives received: {directives}", flush=True)
            return
        try:
            applied = node_ref.apply_tracker_directives(directives)
        except Exception as exc:
            print(f"  [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
            return
        if applied:
            model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
            register_payload["model"] = str(model_id).split("/")[-1]
            register_payload["hf_repo"] = model_id
            register_payload["shard_start"] = applied["shard_start"]
            register_payload["shard_end"] = applied["shard_end"]
            register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization"))
            register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False))

    def _loop() -> None:
        nonlocal node_id
        hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
        outage_streak = 0  # consecutive intervals where tracker was unreachable

        while True:
            time.sleep(interval)

            if outage_streak > 0:
                # Tracker was down — attempt re-registration first (it may have restarted
                # with a clean slate and won't know this node).
                if _reregister():
                    hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
                    print(f"  [node] re-registered after outage — node ID: {node_id}", flush=True)
                    outage_streak = 0
                else:
                    outage_streak += 1
                    if outage_streak <= 3 or outage_streak % 10 == 0:
                        print(
                            f"  [node] WARNING: tracker still unreachable "
                            f"({outage_streak * interval:.0f}s)",
                            flush=True,
                        )
                continue

            try:
                resp = _post_json(hb_url, _get_stats())
                _apply_directives(resp.get("directives", []))
                new_asgn = resp.get("new_assignment")
                if new_asgn:
                    print(
                        f"  [node] tracker reassignment received: "
                        f"model={new_asgn.get('model')!r} "
                        f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
                        flush=True,
                    )
            except urllib.error.HTTPError as exc:
                if exc.code == 404:
                    # Node was purged (e.g. long gap before restart noticed) — re-register now.
                    print("  [node] tracker lost registration — re-registering...", flush=True)
                    if _reregister():
                        hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
                        print(f"  [node] re-registered — node ID: {node_id}", flush=True)
                    else:
                        print("  [node] WARNING: re-registration failed", flush=True)
                        outage_streak = 1
                else:
                    print(f"  [node] WARNING: heartbeat failed ({exc.code}): {exc}", flush=True)
            except Exception as exc:
                outage_streak = 1
                print(f"  [node] WARNING: tracker unreachable: {exc}", flush=True)

    t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
    t.start()
    return t


def _warn_virtual_network_ip(ip: str | None) -> None:
    """Print a warning when the auto-detected advertise IP is in a known virtual-network range.

    172.16.0.0/12 is used by Docker, WSL2, and most hypervisors.  Nodes behind these
    adapters are NOT directly reachable from other physical machines on the LAN, so
    cross-host pipeline hops will time out.  The user must pass --advertise-host with
    their actual LAN IP (e.g. 192.168.x.x) to fix this.
    """
    if ip is None:
        return
    try:
        parts = [int(p) for p in ip.split(".")]
        if len(parts) != 4:
            return
        a, b = parts[0], parts[1]
        # 172.16.0.0/12 → 172.16–31.x.x
        if a == 172 and 16 <= b <= 31:
            print(
                f"\n  WARNING: auto-detected endpoint IP {ip} is in 172.16.0.0/12.\n"
                f"  This range is used by Docker, WSL2, and virtual machines and is\n"
                f"  NOT reachable from other physical machines on your LAN.\n"
                f"  Cross-host pipeline hops WILL time out.\n"
                f"  Fix: use a public tracker with relay (wss://…/ws), or pass\n"
                f"  --advertise-host <your-LAN-ip> (e.g. 192.168.x.x).\n",
                flush=True,
            )
    except Exception:
        pass


def run_startup(
    tracker_url: str,
    port: int = 0,
    model: str = "stub-model",
    model_id: str | None = None,
    shard_start: int | None = None,
    shard_end: int | None = None,
    quantization: str = "bfloat16",
    wallet_path: Path | None = None,
    cache_dir: Path | None = None,
    host: str = "127.0.0.1",
    advertise_host: str | None = None,
    contracts: Any | None = None,
    route_timeout: float = 30.0,
    vram_mb_override: int | None = None,
    max_loaded_shards: int = 1,
    debug: bool = False,
) -> StubNodeServer | TorchNodeServer:
    """Execute the full startup sequence and return a running node server.

    Steps (all non-interactive):
      1. Detect GPU / hardware profile
      2. Load or generate Solana wallet keypair
      3. Query tracker for optimal shard assignment
      4. Download (or stub) the assigned shard from peers, then HuggingFace
      5. Start local HTTP server
      6. Register with tracker

    Prints a compact status summary on completion.
    """

    tracker_url = tracker_url.rstrip("/")
    relay_url = _discover_relay_url(tracker_url)
    if max_loaded_shards < 1:
        raise ValueError("--max-shards must be at least 1")

    # 1. Hardware detection
    if advertise_host is None and host == "0.0.0.0":
        # socket.getfqdn() returns an mDNS name (.local / .localdomain) that remote
        # machines on a different OS or subnet often can't resolve. Instead, probe the
        # outbound IP by opening a UDP socket toward the tracker — no data is sent.
        try:
            _tracker_host = urllib.parse.urlparse(tracker_url).hostname or "8.8.8.8"
            _s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            _s.connect((_tracker_host, 80))
            advertise_host = _s.getsockname()[0]
            _s.close()
        except Exception:
            advertise_host = socket.getfqdn()

    if relay_url:
        print(f"Relay advertised by tracker — using outbound tunnel {relay_url}", flush=True)
    else:
        _warn_virtual_network_ip(advertise_host)

    print("Detecting hardware...", flush=True)
    hw = detect_hardware()
    device: str = hw["device"]
    gpu_name: str | None = hw.get("gpu_name")
    vram_mb: int = hw.get("vram_mb", 0)
    shared_vram_mb: int = hw.get("shared_vram_mb", 0)
    ram_mb: int = hw.get("ram_mb", 16 * 1024)

    if vram_mb_override is not None:
        vram_mb = vram_mb_override
        shared_vram_mb = 0
        print(f"  Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True)
    elif device == "cpu":
        gpu_suffix = ""
        if gpu_name and vram_mb > 0:
            gpu_suffix = (
                f"; CUDA inactive; detected {gpu_name} "
                f"({vram_mb / 1024:.1f} GB dedicated VRAM, {shared_vram_mb / 1024:.1f} GB shared)"
            )
        print(f"  WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM{gpu_suffix})", flush=True)
    else:
        shared_suffix = f", {shared_vram_mb / 1024:.1f} GB shared" if shared_vram_mb > 0 else ""
        print(f"  GPU: {gpu_name} ({vram_mb / 1024:.1f} GB dedicated VRAM{shared_suffix}, {ram_mb / 1024:.1f} GB RAM)", flush=True)

    if vram_mb_override is not None:
        memory_budget_mb = vram_mb
        memory_budget_source = "memory override"
    else:
        memory_budget_mb, memory_budget_source = _memory_budget(device, vram_mb, ram_mb, shared_vram_mb)
    assignment_vram_mb = memory_budget_mb if device == "cuda" or vram_mb_override is not None else 0
    print(f"  Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)

    print("Benchmarking compute...", flush=True)
    bench_tps = benchmark_throughput(device)
    device_label = "GPU" if device == "cuda" else "CPU"
    print(f"  {device_label} throughput index: {bench_tps:,.0f}", flush=True)

    registration_capabilities = {
        "vram_bytes": max(0, int(assignment_vram_mb)) * 1024 * 1024,
        "ram_bytes": max(0, int(ram_mb)) * 1024 * 1024,
        "max_loaded_shards": max_loaded_shards,
        "benchmark_tokens_per_sec": bench_tps,
    }
    # 2. Wallet
    print("Loading wallet...", flush=True)
    wallet_kwargs: dict = {}
    if wallet_path is not None:
        wallet_kwargs["path"] = wallet_path
    _, _, address = load_or_create_wallet(**wallet_kwargs)
    print(f"  Wallet: {address}", flush=True)
    probationary_line = _probationary_status_line(contracts, address)
    if probationary_line is not None:
        print(f"  {probationary_line}", flush=True)

    if model_id:  # treat "" the same as None — no explicit model given
        user_pinned_shard = shard_start is not None or shard_end is not None
        # Auto-detect shard range from model config if not explicitly provided
        if shard_start is None or shard_end is None:
            detected = _detect_num_layers(model_id)
            if detected is None:
                raise ValueError(
                    f"Could not read num_hidden_layers from {model_id} config. "
                    "Pass --shard-start and --shard-end explicitly."
                )
            # When no explicit shard range given, ask the tracker if there's a gap for this model.
            if shard_start is None and shard_end is None:
                try:
                    qs = urllib.parse.urlencode({
                        "device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb, "hf_repo": model_id,
                    })
                    net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
                    if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
                        shard_start = net_asgn["shard_start"]
                        shard_end = net_asgn["shard_end"]
                        print(
                            f"  Tracker found uncovered shard: "
                            f"layers {shard_start}–{shard_end} (of {detected})",
                            flush=True,
                        )
                except Exception:
                    pass  # No other nodes registered yet — default to full model below
            shard_start = shard_start if shard_start is not None else 0
            shard_end = shard_end if shard_end is not None else detected - 1
            print(f"  Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)

        print("Loading real PyTorch model shard...", flush=True)
        node = TorchNodeServer(
            host=host,
            port=port,
            model_id=model_id,
            shard_start=shard_start,
            shard_end=shard_end,
            quantization=quantization,
            tracker_url=tracker_url,
            route_timeout=route_timeout,
            debug=debug,
        )
        _node_start_time = time.monotonic()
        actual_port = node.start()
        total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
        if isinstance(total_layers, int) and total_layers > 0:
            layer_count = shard_end - shard_start + 1
            shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
        else:
            shard_label = f"layers {shard_start}–{shard_end}"
        public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
        endpoint = f"http://{public_host}:{actual_port}"
        local_base_url = f"http://127.0.0.1:{actual_port}"
        relay_bridge, relay_fields = _start_relay_bridge_if_available(
            tracker_url,
            address,
            local_base_url,
            endpoint,
            relay_url=relay_url,
        )
        _attach_relay_bridge(node, relay_bridge)
        # Register with tracker so other nodes can auto-join this model.
        total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
        reg_payload = {
            "endpoint": endpoint,
            "model": model_id.split("/")[-1],
            "hf_repo": model_id,
            "num_layers": total_layers,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "hardware_profile": hw,
            "wallet_address": address,
            "quantization": quantization,
            "score": 1.0,
            "tracker_mode": (shard_start == 0),
            "managed_assignment": not user_pinned_shard,
            **registration_capabilities,
            **relay_fields,
        }
        tracker_node_id: str | None = None
        try:
            reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
            tracker_node_id = str(reg_resp.get("node_id") or "?")
            setattr(node, "tracker_node_id", tracker_node_id)
            print(f"  Registered with tracker — node ID: {tracker_node_id}", flush=True)
            _start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
        except Exception as exc:
            setattr(node, "tracker_node_id", None)
            print(f"  Warning: tracker registration failed: {exc}", flush=True)

        print(
            f"\n{'=' * 32}\n"
            f"meshnet-node ready\n"
            f"  Wallet:       {address}\n"
            f"  Model ID:     {model_id}\n"
            f"  Shard:        {shard_label}\n"
            f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization)}\n"
            f"  Quantization: {quantization}\n"
            f"  Endpoint:     {endpoint}\n"
            f"  Node ID:      {tracker_node_id or 'unregistered'}\n"
            f"  Hardware:     {_hardware_label(device, gpu_name)}\n"
            f"  Benchmark:    {bench_tps:,.0f} (throughput index)\n"
            f"{'=' * 32}",
            flush=True,
        )
        return node
    if shard_start is not None or shard_end is not None:
        raise ValueError("--shard-start / --shard-end require --model-id")

    # 3a. Auto-join: query tracker for network-wide HF model assignment.
    print("Querying tracker for network assignment...", flush=True)
    assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
    net_assignment: dict = {}
    try:
        net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
    except Exception as exc:
        print(f"  (auto-join unavailable: {exc})", flush=True)
    assigned_hf_repo: str | None = net_assignment.get("hf_repo")
    _gap_found: bool = bool(net_assignment.get("gap_found", False))

    if assigned_hf_repo and _gap_found:
        assigned_shard_start: int = net_assignment["shard_start"]
        assigned_shard_end: int = net_assignment["shard_end"]
        assigned_num_layers: int = net_assignment["num_layers"]
        print(
            f"  Assigned: {assigned_hf_repo} "
            f"layers {assigned_shard_start}–{assigned_shard_end} "
            f"(of {assigned_num_layers})",
            flush=True,
        )
        print("Loading real PyTorch model shard...", flush=True)
        node = TorchNodeServer(
            host=host,
            port=port,
            model_id=assigned_hf_repo,
            shard_start=assigned_shard_start,
            shard_end=assigned_shard_end,
            quantization=quantization,
            tracker_url=tracker_url,
            route_timeout=route_timeout,
            debug=debug,
        )
        _node_start_time = time.monotonic()
        actual_port = node.start()
        public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
        endpoint = f"http://{public_host}:{actual_port}"
        local_base_url = f"http://127.0.0.1:{actual_port}"
        relay_bridge, relay_fields = _start_relay_bridge_if_available(
            tracker_url,
            address,
            local_base_url,
            endpoint,
            relay_url=relay_url,
        )
        _attach_relay_bridge(node, relay_bridge)
        auto_reg_payload = {
            "endpoint": endpoint,
            "model": assigned_hf_repo.split("/")[-1],
            "hf_repo": assigned_hf_repo,
            "num_layers": assigned_num_layers,
            "shard_start": assigned_shard_start,
            "shard_end": assigned_shard_end,
            "hardware_profile": hw,
            "wallet_address": address,
            "quantization": quantization,
            "score": 1.0,
            "tracker_mode": (assigned_shard_start == 0),
            "managed_assignment": True,
            **registration_capabilities,
            **relay_fields,
        }
        tracker_node_id = None
        try:
            reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
            tracker_node_id = str(reg_resp.get("node_id") or "?")
            setattr(node, "tracker_node_id", tracker_node_id)
            print(f"  Registered with tracker — node ID: {tracker_node_id}", flush=True)
            _start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
        except Exception as exc:
            setattr(node, "tracker_node_id", None)
            print(f"  Warning: tracker registration failed: {exc}", flush=True)
        shard_count = assigned_shard_end - assigned_shard_start + 1
        print(
            f"\n{'=' * 32}\n"
            f"meshnet-node ready (auto-joined)\n"
            f"  Wallet:       {address}\n"
            f"  Model ID:     {assigned_hf_repo}\n"
            f"  Shard:        layers {assigned_shard_start}–{assigned_shard_end} "
            f"({shard_count} of {assigned_num_layers})\n"
            f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
            f"  Quantization: {quantization}\n"
            f"  Endpoint:     {endpoint}\n"
            f"  Node ID:      {tracker_node_id or 'unregistered'}\n"
            f"  Hardware:     {_hardware_label(device, gpu_name)}\n"
            f"  Benchmark:    {bench_tps:,.0f} (throughput index)\n"
            f"{'=' * 32}",
            flush=True,
        )
        return node

    # 3b. Shard assignment from tracker (stub-model / preset-based path)
    print("Querying tracker for shard assignment...", flush=True)
    assign_qs = urllib.parse.urlencode({
        "model": model,
        "device": device,
        "vram_mb": vram_mb,
        "ram_mb": ram_mb,
    })
    try:
        assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
    except urllib.error.URLError as exc:
        print(f"  ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
        raise

    shard_start: int = assignment["shard_start"]
    shard_end: int = assignment["shard_end"]
    assigned_model: str = assignment.get("model", model)
    hf_repo: str | None = assignment.get("hf_repo")
    peers: list[dict] = assignment.get("peers", [])
    print(f"  Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)

    # 4. Download shard
    print("Downloading shard...", flush=True)
    dl_kwargs: dict = {}
    if cache_dir is not None:
        dl_kwargs["cache_dir"] = cache_dir
    if hf_repo is not None:
        dl_kwargs["hf_repo"] = hf_repo
    if peers:
        dl_kwargs["peers"] = peers
    shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
    shard_checksum = compute_shard_checksum(shard_path)
    print(f"  Cached at: {shard_path}", flush=True)

    # 5. Start HTTP server
    is_last = shard_end >= assignment.get("model_layers_end", shard_end)
    node = StubNodeServer(
        host=host,
        port=port,
        shard_start=shard_start,
        shard_end=shard_end,
        is_last_shard=is_last,
        model=assigned_model,
        shard_path=shard_path,
    )
    actual_port = node.start()
    public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
    endpoint = f"http://{public_host}:{actual_port}"
    local_base_url = f"http://127.0.0.1:{actual_port}"
    relay_bridge, relay_fields = _start_relay_bridge_if_available(
        tracker_url,
        address,
        local_base_url,
        endpoint,
        relay_url=relay_url,
    )
    _attach_relay_bridge(node, relay_bridge)

    # 6. Register with tracker
    print("Registering with tracker...", flush=True)
    try:
        reg_resp = _post_json(
            f"{tracker_url}/v1/nodes/register",
            {
                "endpoint": endpoint,
                "model": assigned_model,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "shard_checksum": shard_checksum,
                "hardware_profile": hw,
                "wallet_address": address,
                "score": 1.0,
                **registration_capabilities,
                **relay_fields,
            },
        )
        node_id = str(reg_resp["node_id"])
        setattr(node, "tracker_node_id", node_id)
    except Exception:
        node.stop()
        raise

    # Status summary
    hw_str = device.upper()
    if gpu_name:
        hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
    print(
        f"\n{'=' * 32}\n"
        f"meshnet-node ready\n"
        f"  Wallet:   {address}\n"
        f"  Shard:    layers {shard_start}-{shard_end} ({assigned_model})\n"
        f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization)}\n"
        f"  Endpoint: {endpoint}\n"
        f"  Node ID:  {node_id}\n"
        f"  Hardware: {hw_str}\n"
        f"  Benchmark: {bench_tps:,.0f} (throughput index)\n"
        f"{'=' * 32}",
        flush=True,
    )

    return node


def _detect_num_layers(model_id: str) -> int | None:
    """Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
    try:
        from transformers import AutoConfig  # type: ignore[import]
        cfg = AutoConfig.from_pretrained(model_id)
        return int(cfg.num_hidden_layers)
    except Exception as exc:
        print(f"  Warning: could not read model config from HF: {exc}", flush=True)
        return None


def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
    if contracts is None:
        return None
    remaining = contracts.registry.probationary_jobs_remaining(wallet_address)
    if remaining <= 0:
        return "Probationary period complete: earning enabled"
    suffix = "job" if remaining == 1 else "jobs"
    return f"Probationary period: {remaining} {suffix} remaining before earning"