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

from __future__ import annotations

import json
import os
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_checked
from .model_catalog import model_metadata_for
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 _downloaded_model_inventory(
    model: str,
    shard_start: int,
    shard_end: int,
    shard_path: Path,
    hf_repo: str | None = None,
    model_sources: list[dict] | None = None,
) -> list[dict]:
    """Return a cheap local inventory record without reading model file contents."""
    file_count = 0
    total_bytes = 0
    existing_rel_files: set[str] = set()
    if shard_path.exists():
        for path in shard_path.rglob("*"):
            if not path.is_file():
                continue
            file_count += 1
            try:
                existing_rel_files.add(path.relative_to(shard_path).as_posix())
            except ValueError:
                pass
            try:
                total_bytes += path.stat().st_size
            except OSError:
                pass
    record = {
        "model": model,
        "shard_start": shard_start,
        "shard_end": shard_end,
        "path": str(shard_path),
        "file_count": file_count,
        "total_bytes": total_bytes,
    }
    if hf_repo is not None:
        record["hf_repo"] = hf_repo
    expected_files: set[str] = set()
    file_sizes: dict[str, int] = {}
    for source in model_sources or []:
        for rel in source.get("full_files") or source.get("files") or []:
            if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
                expected_files.add(rel)
        sizes = source.get("file_sizes")
        if isinstance(sizes, dict):
            for rel, size in sizes.items():
                if isinstance(rel, str) and isinstance(size, int):
                    file_sizes[rel] = size
    if expected_files:
        expected_bytes = sum(file_sizes.get(rel, 0) for rel in expected_files)
        local_expected_files = existing_rel_files & expected_files
        local_expected_bytes = sum(file_sizes.get(rel, 0) for rel in local_expected_files)
        record["expected_file_count"] = len(expected_files)
        record["local_expected_file_count"] = len(local_expected_files)
        record["expected_bytes"] = expected_bytes
        record["local_expected_bytes"] = local_expected_bytes
        record["local_model_percentage"] = (
            round((local_expected_bytes / expected_bytes) * 100, 4)
            if expected_bytes > 0
            else 0.0
        )
    return [record]


def _registration_shard_checksum(model: str, shard_path: Path) -> str | None:
    """Only checksum tiny stub shards; real model folders are too large to hash at startup."""
    if model != "stub-model":
        return None
    return compute_shard_checksum(shard_path)


def _model_cache_path(model_id: str, cache_dir: Path | None) -> Path | None:
    if cache_dir is None:
        return None
    if (cache_dir / "config.json").exists():
        return cache_dir
    candidate = cache_dir / model_id.split("/")[-1]
    if candidate.exists():
        return candidate
    return None


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 _full_model_sources(model_sources: list[dict]) -> list[dict]:
    """Use tracker full-snapshot URLs for real HF model loading."""
    full_sources: list[dict] = []
    for source in model_sources:
        full_url = source.get("full_url")
        if isinstance(full_url, str) and full_url:
            full_sources.append({
                **source,
                "url": full_url,
                # full_files (when advertised) enables robust per-file download
                # of the whole snapshot; empty list falls back to the tar stream.
                "files": source.get("full_files") or [],
                "type": f"{source.get('type') or 'model-source'}-full",
            })
    return full_sources


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 _positive_int(value: int | str | None, name: str) -> int | None:
    if value is None or value == "":
        return None
    try:
        parsed = int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{name} must be a positive integer") from exc
    if parsed < 1:
        raise ValueError(f"{name} must be a positive integer")
    return parsed


def _configure_torch_threads(
    torch_threads: int | None = None,
    torch_interop_threads: int | None = None,
) -> dict[str, int]:
    """Apply PyTorch CPU thread settings before model load/benchmark."""
    intra_threads = _positive_int(
        torch_threads if torch_threads is not None else os.environ.get("MESHNET_TORCH_THREADS"),
        "--torch-threads",
    )
    interop_threads = _positive_int(
        torch_interop_threads
        if torch_interop_threads is not None
        else os.environ.get("MESHNET_TORCH_INTEROP_THREADS"),
        "--torch-interop-threads",
    )

    if intra_threads is not None:
        os.environ.setdefault("OMP_NUM_THREADS", str(intra_threads))
        os.environ.setdefault("MKL_NUM_THREADS", str(intra_threads))
    try:
        import torch
    except ModuleNotFoundError:
        return {}

    if intra_threads is not None:
        torch.set_num_threads(intra_threads)
    if interop_threads is not None:
        torch.set_num_interop_threads(interop_threads)

    active: dict[str, int] = {}
    try:
        active["torch_threads"] = int(torch.get_num_threads())
    except Exception:
        pass
    try:
        active["torch_interop_threads"] = int(torch.get_num_interop_threads())
    except Exception:
        pass
    return active


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


def _runtime_shard_safety_fraction(device: str) -> float:
    """CPU partial loads need room for model skeletons, tokenizer, and allocator peaks."""
    return 0.55 if device != "cuda" else 0.8


def _cap_auto_assigned_shard(
    shard_start: int,
    shard_end: int,
    total_layers: int | None,
    memory_mb: int,
    bytes_per_layer: int | None,
    device: str,
) -> tuple[int, bool, int]:
    if bytes_per_layer is None or total_layers is None:
        return shard_end, False, 0
    max_layers = _max_assignable_layers(
        memory_mb,
        total_layers,
        bytes_per_layer=bytes_per_layer,
        safety_fraction=_runtime_shard_safety_fraction(device),
    )
    if max_layers <= 0:
        return shard_end, False, max_layers
    assigned_layers = shard_end - shard_start + 1
    if assigned_layers <= max_layers:
        return shard_end, False, max_layers
    return min(total_layers - 1, shard_start + max_layers - 1), True, max_layers


def _format_shard_label(
    shard_start: int,
    shard_end: int,
    total_layers: int | None = None,
    *,
    model_name: str | None = None,
) -> str:
    layer_count = shard_end - shard_start + 1
    if isinstance(total_layers, int) and total_layers > 0:
        return f"layers {shard_start}–{shard_end} ({layer_count} of {total_layers})"
    if model_name:
        return f"layers {shard_start}–{shard_end} ({model_name})"
    return f"layers {shard_start}–{shard_end}"


def _shard_budget_line(
    memory_mb: int,
    memory_source: str,
    total_layers: int | None,
    quantization: str,
    bytes_per_layer: int | None = None,
    safety_fraction: float = 0.8,
) -> str:
    memory_gb = memory_mb / 1024
    gb_str = f"{memory_gb:.1f} GB"
    budget_quantization = "bfloat16" if quantization == "auto" else quantization
    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,
        bytes_per_layer=bytes_per_layer,
        safety_fraction=safety_fraction,
    )
    # Remaining capacity after one full model load (rough estimate)
    shard_bytes = max_layers * (bytes_per_layer or _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 {budget_quantization}"
        f"{remaining_str}"
    )


def _assignment_bytes_per_layer(assignment: dict, quantization: str) -> int | None:
    bytes_per_layer = assignment.get("bytes_per_layer")
    if isinstance(bytes_per_layer, int) and bytes_per_layer > 0:
        return bytes_per_layer
    if not isinstance(bytes_per_layer, dict):
        return None
    keys = [quantization, "bfloat16", "bf16", "int8", "nf4"]
    for key in keys:
        value = bytes_per_layer.get(key)
        if isinstance(value, int) and value > 0:
            return value
    for value in bytes_per_layer.values():
        if isinstance(value, int) and value > 0:
            return value
    return None


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]


_PENDING_NODE_ID = "pending"


_HEARTBEAT_INTERVAL_IDLE = 20.0
_HEARTBEAT_INTERVAL_BUSY = 3.0


def _start_heartbeat(
    tracker_url: str,
    node_id: str,
    register_payload: dict,
    interval: float = _HEARTBEAT_INTERVAL_IDLE,
    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, current_requests, 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 _current_requests_snapshot() -> list[dict]:
        if node_ref is None:
            return []
        getter = getattr(node_ref, "current_requests", None)
        if getter is None:
            return []
        current = getter() if callable(getter) else getter
        return list(current) if isinstance(current, list) else []

    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)
        current_requests = _current_requests_snapshot()
        if current_requests:
            stats["current_requests"] = current_requests
        return stats

    def _sleep_interval() -> float:
        if _current_requests_snapshot() or (node_ref is not None and getattr(node_ref, "queue_depth", 0) > 0):
            return _HEARTBEAT_INTERVAL_BUSY
        return interval

    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)
            if node_ref is not None:
                setattr(node_ref, "tracker_node_id", node_id)
            return True
        except Exception:
            return False

    def _register_additional_assignment(applied: dict) -> None:
        model_id = str(applied.get("model") or register_payload.get("hf_repo") or register_payload.get("model"))
        extra_payload = {
            **register_payload,
            "model": model_id.split("/")[-1],
            "hf_repo": model_id if "/" in model_id else register_payload.get("hf_repo"),
            "shard_start": applied["shard_start"],
            "shard_end": applied["shard_end"],
            "quantization": applied.get("quantization", register_payload.get("quantization")),
            "tracker_mode": bool(applied.get("tracker_mode", False)),
            "managed_assignment": True,
        }
        try:
            reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", extra_payload)
            print(
                f"  [node] registered additional model — node ID: {reg_resp.get('node_id')}",
                flush=True,
            )
        except Exception as exc:
            print(f"  [node] WARNING: additional model registration failed: {exc}", flush=True)

    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:
            if applied.get("action") == "ADD_SHARD":
                _register_additional_assignment(applied)
                return
            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 = 1 if node_id == _PENDING_NODE_ID else 0

        while True:
            time.sleep(_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 assignment received: "
                        f"action={new_asgn.get('action')!r} model={new_asgn.get('model')!r} "
                        f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
                        flush=True,
                    )
                    _apply_directives([new_asgn])
            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 _register_with_tracker(
    tracker_url: str,
    reg_payload: dict,
    node: Any,
    start_time: float,
) -> str | None:
    """Register with the tracker, or start background retries when it is unreachable."""
    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=start_time)
        return tracker_node_id
    except Exception as exc:
        setattr(node, "tracker_node_id", None)
        print(f"  Warning: tracker registration failed: {exc}", flush=True)
        print("  [node] will retry registration in the background", flush=True)
        _start_heartbeat(
            tracker_url,
            _PENDING_NODE_ID,
            reg_payload,
            node_ref=node,
            start_time=start_time,
        )
        return None


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 | None = None,
    model_id: str | None = None,
    shard_start: int | None = None,
    shard_end: int | None = None,
    quantization: str = "auto",
    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,
    tracker_source_disabled: bool = False,
    torch_threads: int | None = None,
    torch_interop_threads: int | None = None,
) -> 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()
    torch_thread_config = _configure_torch_threads(torch_threads, torch_interop_threads)
    if torch_thread_config:
        hw.update(torch_thread_config)
        intra = torch_thread_config.get("torch_threads", "?")
        interop = torch_thread_config.get("torch_interop_threads", "?")
        print(f"  PyTorch threads: intra-op={intra}, inter-op={interop}", flush=True)
    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)
    if device != "cuda" and gpu_name:
        _cuda_score, cuda_ok, cuda_error = benchmark_throughput_checked("cuda")
        hw["cuda_benchmark_ok"] = cuda_ok
        if cuda_error:
            hw["cuda_benchmark_error"] = cuda_error
        if not cuda_ok:
            print(f"  CUDA benchmark unavailable: {cuda_error}; using CPU benchmark", flush=True)
    bench_tps, bench_ok, bench_error = benchmark_throughput_checked(device)
    hw["benchmark_device"] = device
    hw["benchmark_ok"] = bench_ok
    if bench_error:
        hw["benchmark_error"] = bench_error
    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)

    pinned_shard_start = shard_start
    pinned_shard_end = shard_end
    user_pinned_shard = pinned_shard_start is not None or pinned_shard_end is not None

    if model_id:  # treat "" the same as None — no explicit model given
        full_sources: list[dict] = []
        # Auto-detect shard range from model config if not explicitly provided
        if shard_start is None or shard_end is None:
            try:
                detected = _detect_num_layers(model_id, cache_dir=cache_dir)
            except TypeError:
                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"]
                        asgn_total_layers = int(net_asgn.get("num_layers") or detected)
                        asgn_bytes_per_layer = _assignment_bytes_per_layer(net_asgn, quantization)
                        capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
                            shard_start,
                            shard_end,
                            asgn_total_layers,
                            memory_budget_mb,
                            asgn_bytes_per_layer,
                            device,
                        )
                        if was_capped:
                            original_end = shard_end
                            shard_end = capped_shard_end
                            print(
                                f"  WARNING: tracker assigned layers {shard_start}-{original_end}, "
                                f"but CPU-safe runtime budget fits {max_runtime_layers}/{asgn_total_layers} layers; "
                                f"loading layers {shard_start}-{shard_end} instead.",
                                flush=True,
                            )
                        full_sources = (
                            [] if tracker_source_disabled
                            else _full_model_sources(net_asgn.get("model_sources", []))
                        )
                        if full_sources:
                            cache_dir = download_shard(
                                model_id.split("/")[-1],
                                shard_start,
                                shard_end,
                                cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
                                hf_repo=model_id,
                                model_sources=full_sources,
                            )
                        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,
            cache_dir=cache_dir,
            debug=debug,
            max_loaded_shards=max_loaded_shards,
        )
        _node_start_time = time.monotonic()
        actual_port = node.start()
        total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
        shard_label = _format_shard_label(shard_start, shard_end, total_layers)
        public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
        endpoint = f"http://{public_host}:{actual_port}"
        node.set_advertised_endpoint(endpoint)
        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)
        model_cache_path = _model_cache_path(model_id, cache_dir)
        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,
            "model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
            "downloaded_models": (
                _downloaded_model_inventory(
                    model_id.split("/")[-1],
                    shard_start,
                    shard_end,
                    model_cache_path,
                    hf_repo=model_id,
                    model_sources=full_sources,
                )
                if model_cache_path is not None
                else []
            ),
            **registration_capabilities,
            **relay_fields,
        }
        tracker_node_id = _register_with_tracker(
            tracker_url, reg_payload, node, _node_start_time,
        )

        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 user_pinned_shard and not model:
        raise ValueError("--shard-start / --shard-end require --model")

    # 3a. Auto-join: query tracker for network-wide HF model assignment.
    # Skipped when the user explicitly requested a model — the shard-assignment
    # query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
    # tracker would otherwise print a scary 503 for the model-less auto-join.
    net_assignment: dict = {}
    if model_id or (model and model != "stub-model"):
        if model:
            print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
    else:
        print("Querying tracker for network assignment...", flush=True)
        assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
        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:
        assigned_shard_start: int = net_assignment["shard_start"]
        assigned_shard_end: int = net_assignment["shard_end"]
        assigned_num_layers: int = net_assignment["num_layers"]
        assigned_bytes_per_layer = _assignment_bytes_per_layer(net_assignment, quantization)
        capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
            assigned_shard_start,
            assigned_shard_end,
            assigned_num_layers,
            memory_budget_mb,
            assigned_bytes_per_layer,
            device,
        )
        if was_capped:
            original_end = assigned_shard_end
            assigned_shard_end = capped_shard_end
            print(
                f"  WARNING: tracker assigned layers {assigned_shard_start}-{original_end}, "
                f"but CPU-safe runtime budget fits {max_runtime_layers}/{assigned_num_layers} layers; "
                f"loading layers {assigned_shard_start}-{assigned_shard_end} instead.",
                flush=True,
            )
        assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
        if _gap_found:
            print(
                f"  Assigned gap: {assigned_hf_repo} "
                f"layers {assigned_shard_start}-{assigned_shard_end} "
                f"(of {assigned_num_layers})",
                flush=True,
            )
        else:
            print(
                f"  Assigned redundant copy: {assigned_hf_repo} "
                f"layers {assigned_shard_start}-{assigned_shard_end} "
                f"(of {assigned_num_layers})",
                flush=True,
            )
        full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
        if full_sources:
            print("Downloading assigned model snapshot...", flush=True)
            cache_dir = download_shard(
                assigned_hf_repo.split("/")[-1],
                assigned_shard_start,
                assigned_shard_end,
                cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
                hf_repo=assigned_hf_repo,
                model_sources=full_sources,
            )
        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,
            cache_dir=cache_dir,
            debug=debug,
            max_loaded_shards=max_loaded_shards,
        )
        _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}"
        node.set_advertised_endpoint(endpoint)
        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)
        model_cache_path = _model_cache_path(assigned_hf_repo, cache_dir)
        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,
            "model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
            "downloaded_models": (
                _downloaded_model_inventory(
                    assigned_hf_repo.split("/")[-1],
                    assigned_shard_start,
                    assigned_shard_end,
                    model_cache_path,
                    hf_repo=assigned_hf_repo,
                    model_sources=full_sources,
                )
                if model_cache_path is not None
                else []
            ),
            **registration_capabilities,
            **relay_fields,
        }
        tracker_node_id = _register_with_tracker(
            tracker_url, auto_reg_payload, node, _node_start_time,
        )
        shard_label = _format_shard_label(
            assigned_shard_start,
            assigned_shard_end,
            assigned_num_layers,
        )
        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:        {shard_label}\n"
            f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization, bytes_per_layer=assigned_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\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 not assigned_hf_repo and model is None:
        raise RuntimeError(
            "Tracker did not assign a model. Join a network that already serves one, "
            "or start with --model <HF_REPO>."
        )

    # 3b. Stub preset path (tests / explicit stub-model) or named preset models.
    print("Querying tracker for shard assignment...", flush=True)
    assign_qs = urllib.parse.urlencode({
        "model": model or "stub-model",
        "device": device,
        # CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
        # VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
        # 79 GB box whose Torch has no CUDA).
        "vram_mb": assignment_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 = assignment["shard_start"]
    shard_end = assignment["shard_end"]
    if user_pinned_shard:
        if pinned_shard_start is not None:
            shard_start = pinned_shard_start
        if pinned_shard_end is not None:
            shard_end = pinned_shard_end
    assigned_model: str = assignment.get("model", model)
    hf_repo: str | None = assignment.get("hf_repo")
    peers: list[dict] = assignment.get("peers", [])
    model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
    assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
    model_layers_end = assignment.get("model_layers_end")
    assigned_total_layers = (
        int(model_layers_end) + 1
        if model_layers_end is not None
        else None
    )
    shard_label = _format_shard_label(
        shard_start,
        shard_end,
        assigned_total_layers,
        model_name=assigned_model,
    )
    if user_pinned_shard:
        shard_label = f"{shard_label} (pinned)"
    if user_pinned_shard and assigned_total_layers and assignment_bytes_per_layer:
        pinned_layers = shard_end - shard_start + 1
        max_layers = _max_assignable_layers(
            memory_budget_mb,
            assigned_total_layers,
            assignment_bytes_per_layer,
            safety_fraction=_runtime_shard_safety_fraction(device),
        )
        if pinned_layers > max_layers:
            raise ValueError(
                f"Pinned shard layers {shard_start}–{shard_end} ({pinned_layers} layers) exceed "
                f"the {memory_budget_mb / 1024:.1f} GB {memory_budget_source} budget "
                f"(fits up to {max_layers}/{assigned_total_layers} layers at bfloat16). "
                "Drop --shard-start/--shard-end to let the tracker auto-assign, or pin a smaller range."
            )
    print(f"  Shard: {shard_label}", 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
    if model_sources:
        dl_kwargs["model_sources"] = model_sources
    shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
    shard_checksum = _registration_shard_checksum(assigned_model, shard_path)
    downloaded_models = _downloaded_model_inventory(
        assigned_model,
        shard_start,
        shard_end,
        shard_path,
        hf_repo=hf_repo,
        model_sources=model_sources,
    )
    print(f"  Cached at: {shard_path}", flush=True)

    # 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub.
    _node_start_time = time.monotonic()
    if hf_repo and assigned_model != "stub-model":
        print("Loading real PyTorch model shard...", flush=True)
        node = TorchNodeServer(
            host=host,
            port=port,
            model_id=hf_repo,
            shard_start=shard_start,
            shard_end=shard_end,
            quantization=quantization,
            tracker_url=tracker_url,
            route_timeout=route_timeout,
            cache_dir=shard_path,
            debug=debug,
            max_loaded_shards=max_loaded_shards,
        )
        actual_port = node.start()
        total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers
        shard_label = _format_shard_label(shard_start, shard_end, total_layers, model_name=assigned_model)
        if user_pinned_shard:
            shard_label = f"{shard_label} (pinned)"
        public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
        endpoint = f"http://{public_host}:{actual_port}"
        node.set_advertised_endpoint(endpoint)
        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)
        reg_payload = {
            "endpoint": endpoint,
            "model": assigned_model,
            "hf_repo": hf_repo,
            "num_layers": total_layers,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "downloaded_models": downloaded_models,
            "hardware_profile": hw,
            "wallet_address": address,
            "quantization": quantization,
            "score": 1.0,
            "tracker_mode": (shard_start == 0),
            "managed_assignment": not user_pinned_shard,
            "model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
            **registration_capabilities,
            **relay_fields,
        }
        tracker_node_id = _register_with_tracker(
            tracker_url, reg_payload, node, _node_start_time,
        )
        print(
            f"\n{'=' * 32}\n"
            f"meshnet-node ready\n"
            f"  Wallet:       {address}\n"
            f"  Model ID:     {hf_repo}\n"
            f"  Shard:        {shard_label}\n"
            f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\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

    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)
    reg_payload = {
        "endpoint": endpoint,
        "model": assigned_model,
        "shard_start": shard_start,
        "shard_end": shard_end,
        "shard_checksum": shard_checksum,
        "downloaded_models": downloaded_models,
        "hardware_profile": hw,
        "wallet_address": address,
        "score": 1.0,
        "managed_assignment": not user_pinned_shard,
        **registration_capabilities,
        **relay_fields,
    }
    try:
        reg_resp = _post_json(
            f"{tracker_url}/v1/nodes/register",
            reg_payload,
        )
        node_id = str(reg_resp["node_id"])
        setattr(node, "tracker_node_id", node_id)
        _start_heartbeat(tracker_url, node_id, reg_payload, node_ref=node, start_time=_node_start_time)
    except Exception:
        node.stop()
        raise

    # Status summary
    hw_str = device.upper()
    if gpu_name:
        hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
    shard_label = _format_shard_label(
        shard_start,
        shard_end,
        assigned_total_layers,
        model_name=assigned_model,
    )
    print(
        f"\n{'=' * 32}\n"
        f"meshnet-node ready\n"
        f"  Wallet:   {address}\n"
        f"  Shard:    {shard_label}\n"
        f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\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, cache_dir: Path | None = None) -> int | None:
    """Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
    try:
        from transformers import AutoConfig  # type: ignore[import]

        from .model_catalog import layers_from_config

        local_model = _model_cache_path(model_id, cache_dir)
        load_source = str(local_model) if local_model is not None else model_id
        cfg = AutoConfig.from_pretrained(
            load_source,
            cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
        )
        layers = layers_from_config(cfg)
        if layers is None:
            print(
                f"  Warning: no layer count in {model_id} config "
                "(checked top level and text_config)", flush=True,
            )
        return 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"