"""Distributed GGUF shard load integration test.

Downloads a small dense-Llama GGUF (TinyLlama 1.1B Q4_K_M ~670 MB),
loads it in shard ranges via the meshnet-range-loader C wrapper, registers
each shard with a live TrackerServer, and verifies routing, range reporting,
and memory scaling.

Set MESHNET_ENABLE_REAL_INFERENCE_TESTS=1 to run.

Evidence class: real-model integration. Downloads ~670 MB on first run.
"""

from __future__ import annotations

import json
import os
import pathlib
import subprocess
import sys
import urllib.error
import urllib.request

import pytest

ROOT = pathlib.Path(__file__).resolve().parent.parent
sys.path[:0] = [
    str(ROOT / "packages" / "tracker"),
    str(ROOT / "packages" / "node"),
    str(ROOT / "packages" / "contracts"),
]

from meshnet_tracker.server import TrackerServer  # noqa: E402 — sys.path prepended above

# Only run when explicitly enabled
pytestmark = pytest.mark.skipif(
    "MESHNET_ENABLE_REAL_INFERENCE_TESTS" not in os.environ,
    reason="set MESHNET_ENABLE_REAL_INFERENCE_TESTS=1 to download and load a real GGUF",
)

# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------

HF_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
GGUF_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
MODEL_URL = f"https://huggingface.co/{HF_REPO}/resolve/main/{GGUF_FILE}"
EXPECTED_LAYERS = 22
GGUF_FLAVOR = GGUF_FILE.replace(".gguf", "")

CACHE_DIR = ROOT / ".cache" / "gguf-models"
LOADER = ROOT / "build" / "dgr-004-final" / "build" / "bin" / "meshnet-range-loader"
LLAMA_LIB_DIR = LOADER.parent


# ---------------------------------------------------------------------------
# Loading helper
# ---------------------------------------------------------------------------


def load_shard(gguf_path: str, start: int, end: int) -> dict:
    """Load a shard range of the GGUF via the C wrapper and return JSON report."""
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = str(LLAMA_LIB_DIR)

    result = subprocess.run(
        [str(LOADER), gguf_path, str(start), str(end)],
        capture_output=True, text=True, timeout=120, env=env,
    )
    if result.returncode != 0:
        # Parse JSON from stdout even on error if it exists
        if result.stdout.strip():
            try:
                return json.loads(result.stdout)
            except json.JSONDecodeError:
                pass
        raise RuntimeError(
            f"meshnet-range-loader [{start}, {end}) failed (exit {result.returncode}):\n"
            f"stderr: {result.stderr[:500]}"
        )

    return json.loads(result.stdout)


# ---------------------------------------------------------------------------
# Download and cache
# ---------------------------------------------------------------------------


def _ensure_model() -> pathlib.Path:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    model_path = CACHE_DIR / GGUF_FILE

    if model_path.exists() and model_path.stat().st_size > 600 * 1024 * 1024:
        return model_path

    print(f"Downloading {MODEL_URL} (~670 MB)...", file=sys.stderr)
    urllib.request.urlretrieve(MODEL_URL, model_path)
    actual_mb = model_path.stat().st_size / (1024 * 1024)
    print(f"Downloaded {GGUF_FLAVOR}: {actual_mb:.0f} MB", file=sys.stderr)
    return model_path


# ---------------------------------------------------------------------------
# Tracker helpers
# ---------------------------------------------------------------------------


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


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


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_whole_model_load_and_report():
    """Load the full TinyLlama GGUF and verify metadata."""
    gguf = _ensure_model()
    report = load_shard(str(gguf), 0, EXPECTED_LAYERS)
    assert report["ok"] is True
    assert report["start_layer"] == 0
    assert report["end_layer"] == EXPECTED_LAYERS
    assert report["mapped_bytes"] > 0
    assert report["resident_bytes"] >= report["mapped_bytes"]


def test_head_shard_is_less_than_full_model():
    """A head-only shard maps fewer bytes than the full model."""
    gguf = _ensure_model()
    head = load_shard(str(gguf), 0, 8)
    full = load_shard(str(gguf), 0, EXPECTED_LAYERS)
    assert head["mapped_bytes"] <= full["mapped_bytes"]


def test_tail_shard_maps_fewer_bytes_than_middle():
    """Fewer layers = fewer bytes (tail has 6 layers, middle has 8)."""
    gguf = _ensure_model()
    middle = load_shard(str(gguf), 8, 16)
    tail = load_shard(str(gguf), 16, EXPECTED_LAYERS)
    # Middle (8 layers) must map more than tail (6 layers)
    assert middle["mapped_bytes"] > tail["mapped_bytes"]


def test_memory_scales_with_owned_range():
    """More layers = more resident bytes."""
    gguf = _ensure_model()
    small = load_shard(str(gguf), 0, 4)
    large = load_shard(str(gguf), 0, 12)
    assert large["mapped_bytes"] > small["mapped_bytes"]
    assert large["resident_bytes"] > small["resident_bytes"]


def test_invalid_range_rejected():
    """Loading a range outside GGUF layer bounds fails closed."""
    gguf = _ensure_model()
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = str(LLAMA_LIB_DIR)
    result = subprocess.run(
        [str(LOADER), str(gguf), str(EXPECTED_LAYERS + 1), str(EXPECTED_LAYERS + 5)],
        capture_output=True, text=True, timeout=30, env=env,
    )
    # Should fail with exit code 1 and an error message
    assert result.returncode != 0
    assert "error" in result.stderr.lower() or "ERROR" in result.stderr


def test_tracker_registers_shard_nodes():
    """Start a tracker, register multiple shard nodes, verify the console."""
    gguf = _ensure_model()

    tracker = TrackerServer(heartbeat_timeout=60.0)
    tracker_port = tracker.start()

    shards = [
        ("head", 0, 8),
        ("middle", 8, 16),
        ("tail", 16, EXPECTED_LAYERS),
    ]
    registered_ids = []
    try:
        for name_tag, start, end in shards:
            report = load_shard(str(gguf), start, end)
            node_id = f"{GGUF_FLAVOR}-{name_tag}"
            registered_ids.append(node_id)

            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {
                    "node_id": node_id,
                    "endpoint": f"http://localhost:{10000 + start}",
                    "model": GGUF_FLAVOR,
                    "num_layers": report["end_layer"] - report["start_layer"],
                    "shard_start": report["start_layer"],
                    "shard_end": report["end_layer"],
                    "hardware_profile": {
                        "mapped_bytes": report["mapped_bytes"],
                        "resident_bytes": report["resident_bytes"],
                    },
                    "score": 1.0,
                },
            )

        # Verify console shows all three registered nodes
        console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
        registered_eps = {
            f"http://localhost:{10000 + start}" for _, start, _ in shards
        }
        found_eps = set()
        for event in console.get("events", []):
            if event.get("message") == "node registered":
                ep = event.get("fields", {}).get("endpoint", "")
                if ep in registered_eps:
                    found_eps.add(ep)
        for _, start, _ in shards:
            expected_ep = f"http://localhost:{10000 + start}"
            assert expected_ep in found_eps, \
                f"endpoint {expected_ep} not found in registration events"

    finally:
        tracker.stop()