"""ADR-0021: dynamic bandit-style route selection with learned statistics."""

import http.server
import json
import random
import threading
import types
import urllib.request

from meshnet_tracker.routing_stats import (
    RouteCandidate,
    RouteStatsStore,
    RoutingConfig,
    choose_route,
    route_signature,
    route_table,
)
from meshnet_tracker.capability import absent_state
from meshnet_tracker.server import TrackerServer, _enumerate_routes


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


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


def _fake_node(node_id, shard_start, shard_end, benchmark=100.0, endpoint=None):
    return types.SimpleNamespace(
        node_id=node_id,
        endpoint=endpoint or f"http://{node_id}:7000",
        model="qwen3.6-35b-a3b",
        hf_repo="unsloth/Qwen3.6-35B-A3B",
        shard_start=shard_start,
        shard_end=shard_end,
        num_layers=40,
        benchmark_tokens_per_sec=benchmark,
        model_tokens_per_sec={},
        queue_depth=0,
        proxy_inflight=0,
        wallet_address=None,
        relay_addr=None,
        # A pre-capability node (NCA-004): routable only under the `compat`
        # policy, which is what these route-enumeration tests exercise.
        capability=absent_state(),
    )


# ---- RouteStatsStore ----------------------------------------------------


def test_route_stats_sample_becomes_proven_and_decays():
    "Route stats sample becomes proven and decays\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=100.0))
    sig = "m|a[0-39]"
    assert store.snapshot(sig, "m", now=0.0)["status"] == "unsampled"
    assert store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
    snap = store.snapshot(sig, "m", now=1.0)
    assert snap["status"] == "proven"
    assert snap["tps"] == 10.0
    # After many half-lives the sample mass decays below the proven threshold.
    assert store.snapshot(sig, "m", now=1000.0)["status"] == "decayed"


def test_route_stats_rejects_near_empty_samples():
    "Route stats rejects near empty samples\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(min_sample_tokens=8))
    assert not store.record_sample("m", "sig", tokens=3, elapsed_seconds=1.0)
    assert store.snapshot("sig", "m")["samples"] == 0


def test_route_stats_epoch_bump_marks_stale():
    "Route stats epoch bump marks stale\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore()
    sig = "m|a[0-39]"
    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
    assert store.snapshot(sig, "m", now=1.0)["status"] == "proven"
    store.bump_epoch(["m"])
    snap = store.snapshot(sig, "m", now=1.0)
    assert snap["status"] == "stale"
    assert snap["tps"] == 10.0  # EWMA kept as a prior for display
    # A fresh sample under the new epoch re-proves the route.
    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=2.0)
    assert store.snapshot(sig, "m", now=3.0)["status"] == "proven"


def test_route_stats_ewma_averages_samples():
    "Route stats ewma averages samples\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
    sig = "m|a"
    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)  # 10 tps
    store.record_sample("m", sig, tokens=200, elapsed_seconds=10.0, now=1.0)  # 20 tps
    snap = store.snapshot(sig, "m", now=2.0)
    assert 14.9 < snap["tps"] < 15.1


# ---- choose_route --------------------------------------------------------


def test_route_stats_persist_historical_hop_latency_across_restart(tmp_path):
    "Route stats persist historical hop latency across restart\n\nTags: http, performance, persistence, routing, tracker"
    db_path = str(tmp_path / "routing-stats.sqlite")
    signature = "m|gpu[0-11]->cpu[12-23]"
    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9), db_path=db_path)
    store.record_sample("m", signature, tokens=100, elapsed_seconds=2.0, now=100.0)
    store.save_to_db()

    restored = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9), db_path=db_path)
    row = restored.model_rows("m", now=101.0)[0]
    assert row["hop_count"] == 2
    assert row["samples"] == 1
    assert row["tps"] == 50.0
    assert row["latency_ms"] == 2000.0


def _candidates_two_routes():
    fast = RouteCandidate(nodes=[], signature="m|fast", prior_tps=100.0)
    slow = RouteCandidate(nodes=[], signature="m|slow", prior_tps=50.0)
    return fast, slow


def test_choose_route_without_samples_is_deterministic_best_prior():
    "Choose route without samples is deterministic best prior\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore()
    fast, slow = _candidates_two_routes()
    for _ in range(20):
        picked, decision = choose_route([slow, fast], store, "m", rng=random.Random(7))
        assert picked is fast
        assert decision["mode"] == "scout"


def test_choose_route_traffic_proportional_to_tps():
    "Choose route traffic proportional to tps\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
    fast, slow = _candidates_two_routes()
    now = 0.0
    for _ in range(5):
        now += 1.0
        store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
        store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
    rng = random.Random(42)
    picks = {"m|fast": 0, "m|slow": 0}
    for _ in range(4000):
        picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
        assert decision["mode"] == "exploit"
        picks[picked.signature] += 1
    share = picks["m|fast"] / 4000
    # 15 tps vs 10 tps at alpha=1 → expected fast share 0.6
    assert 0.55 < share < 0.65


def test_choose_route_scouts_unproven_routes_at_explore_share():
    "Choose route scouts unproven routes at explore share\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(explore_share=0.25, stats_half_life_seconds=1e9))
    fast, slow = _candidates_two_routes()
    now = 1.0
    store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
    rng = random.Random(11)
    scouted = 0
    for _ in range(4000):
        picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
        if decision["mode"] == "scout":
            scouted += 1
            assert picked is slow
    assert 0.20 < scouted / 4000 < 0.30


# ---- _enumerate_routes ---------------------------------------------------


def test_enumerate_routes_mixed_topology_yields_both_routes():
    "Enumerate routes mixed topology yields both routes\n\nTags: http, performance, routing, tracker"
    gpu = _fake_node("gpu", 0, 21, benchmark=11000.0)
    cpu = _fake_node("cpu", 0, 39, benchmark=425.0)
    candidates = _enumerate_routes([gpu, cpu], 0, 39, model="qwen3.6-35b-a3b")
    signatures = {c.signature for c in candidates}
    assert signatures == {
        route_signature("qwen3.6-35b-a3b", [gpu, cpu]),
        route_signature("qwen3.6-35b-a3b", [cpu]),
    }
    hybrid = next(c for c in candidates if len(c.nodes) == 2)
    assert [n.node_id for n in hybrid.nodes] == ["gpu", "cpu"]
    # Hybrid route's prior is its bottleneck hop, not the fast head.
    assert hybrid.prior_tps == 425.0


def test_enumerate_routes_requires_head_at_first_layer():
    "Enumerate routes requires head at first layer\n\nTags: http, performance, routing, tracker"
    tail_only = _fake_node("tail", 22, 39)
    assert _enumerate_routes([tail_only], 0, 39, model="m") == []


def test_route_table_reports_coefficient_and_share():
    "Route table reports coefficient and share\n\nTags: http, performance, routing, tracker"
    store = RouteStatsStore(RoutingConfig(explore_share=0.3, stats_half_life_seconds=1e9))
    fast, slow = _candidates_two_routes()
    now = 1.0
    for _ in range(3):
        store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
        store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
        now += 1.0
    rows = route_table([fast, slow], store, "m", now=now)
    by_sig = {r["signature"]: r for r in rows}
    assert by_sig["m|fast"]["coefficient"] == 1.0
    assert abs(by_sig["m|slow"]["coefficient"] - (10.0 / 15.0)) < 0.01
    # No scouts → full exploit budget split 0.6 / 0.4.
    assert abs(by_sig["m|fast"]["expected_share"] - 0.6) < 0.01
    assert abs(by_sig["m|slow"]["expected_share"] - 0.4) < 0.01


# ---- integration: proxy uses route head + /v1/routing --------------------


def test_proxy_head_is_route_head_and_routing_endpoint_lists_routes():
    "Mixed topology (partial head 0-21 + full node 0-39): the proxy target must be the selected route's own head, downstream hops must continue at head.\n\nTags: http, performance, routing, tracker"


    class ChatHandler(http.server.BaseHTTPRequestHandler):
        def log_message(self, *args):  # noqa: ARG002
            pass

        def do_POST(self):
            length = int(self.headers.get("Content-Length", 0))
            self.rfile.read(length)
            route_header = self.headers.get("X-Meshnet-Route") or "[]"
            body = json.dumps({
                "choices": [{"message": {"role": "assistant", "content": route_header}}],
                "usage": {"prompt_tokens": 10, "completion_tokens": 40},
            }).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

    stubs = []
    threads = []
    for _ in range(2):
        stub = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
        thread = threading.Thread(target=stub.serve_forever, daemon=True)
        thread.start()
        stubs.append(stub)
        threads.append(thread)
    gpu_stub, cpu_stub = stubs

    tracker = TrackerServer(enable_billing=False, model_presets={
        "qwen3.6-35b-a3b": {
            "layers_start": 0,
            "layers_end": 39,
            "hf_repo": "unsloth/Qwen3.6-35B-A3B",
            "aliases": ["Qwen3.6-35B-A3B"],
        }
    })
    tracker_port = tracker.start()
    try:
        tracker._server.route_rng = random.Random(3)
        for stub, shard_end, bench in ((gpu_stub, 21, 11000.0), (cpu_stub, 39, 425.0)):
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {"endpoint": f"http://127.0.0.1:{stub.server_address[1]}",
                 "model": "qwen3.6-35b-a3b",
                 "hf_repo": "unsloth/Qwen3.6-35B-A3B",
                 "num_layers": 40,
                 "shard_start": 0,
                 "shard_end": shard_end,
                 "tracker_mode": True,
                 "quantization": "bfloat16",
                 "benchmark_tokens_per_sec": bench,
                 "hardware_profile": {},
                 "score": 1.0},
            )

        for _ in range(8):
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
                {"model": "Qwen3.6-35B-A3B",
                 "messages": [{"role": "user", "content": "hi"}]},
            )

        console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
        routing = _get_json(f"http://127.0.0.1:{tracker_port}/v1/routing")
    finally:
        tracker.stop()
        for stub, thread in zip(stubs, threads):
            stub.shutdown()
            stub.server_close()
            thread.join(timeout=1.0)

    gpu_endpoint = f"http://127.0.0.1:{gpu_stub.server_address[1]}"
    cpu_endpoint = f"http://127.0.0.1:{cpu_stub.server_address[1]}"
    selected = [e for e in console["events"] if e["message"] == "proxy route selected"]
    assert selected
    for event in selected:
        fields = event["fields"]
        nodes = fields["nodes"]
        # The proxy head must be the route's first hop (ADR-0020 regression).
        assert fields["head_endpoint"] == nodes[0]["endpoint"]
        downstream = json.loads(fields["downstream"])
        if fields["head_endpoint"] == gpu_endpoint:
            # Partial head: downstream continues at layer 22, never 0.
            assert downstream == [{"endpoint": cpu_endpoint, "start_layer": 22}]
        else:
            assert fields["head_endpoint"] == cpu_endpoint
            assert downstream == []

    table = routing["models"]["qwen3.6-35b-a3b"]
    assert len(table["routes"]) == 2
    sampled = [r for r in table["routes"] if r["samples"] > 0]
    assert sampled, "completed requests must produce route samples"


def test_admin_model_load_request_queues_directive_on_joined_node():
    "Admin model load request queues directive on joined node\n\nTags: http, performance, routing, tracker"
    tracker = TrackerServer(enable_billing=False, validator_service_token="test-admin")
    port = tracker.start()
    try:
        node = _post_json(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            {
                "endpoint": "http://127.0.0.1:9911",
                "model": "stub-model",
                "shard_start": 0,
                "shard_end": 3,
                "managed_assignment": True,
                "memory_mb": 32768,
                "hardware_profile": {"host_id": "available-ram-pool"},
            },
        )
        request = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/models/load",
            data=json.dumps({"model": "qwen2.5-0.5b-instruct"}).encode(),
            headers={"Content-Type": "application/json", "Authorization": "Bearer test-admin"},
            method="POST",
        )
        with urllib.request.urlopen(request) as response:
            result = json.loads(response.read())
        heartbeat = _post_json(
            f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat", {})
    finally:
        tracker.stop()

    assert result["status"] == "queued"
    assert result["assignment"]["node_id"] == node["node_id"]
    assert heartbeat["directives"][0]["action"] == "ADD_SHARD"
    assert heartbeat["directives"][0]["model"] == "Qwen/Qwen2.5-0.5B-Instruct"


def test_admin_can_replace_a_served_model_and_release_it():
    "Forced admin placement replaces a served shard; release queues DROP_SHARD."
    tracker = TrackerServer(enable_billing=False, validator_service_token="test-admin")
    port = tracker.start()
    try:
        node = _post_json(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9912", "model": "stub-model",
             "shard_start": 0, "shard_end": 3, "managed_assignment": True,
             "max_loaded_shards": 1, "memory_mb": 1,
             "hardware_profile": {"host_id": "full-host"}},
        )
        headers = {"Content-Type": "application/json", "Authorization": "Bearer test-admin"}
        load = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/models/load",
            data=json.dumps({
                "model": "qwen2.5-0.5b-instruct",
                "node_id": node["node_id"],
                "force": True,
            }).encode(),
            headers=headers, method="POST")
        with urllib.request.urlopen(load) as response:
            assert json.loads(response.read())["assignment"]["node_id"] == node["node_id"]
        heartbeat = _post_json(f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat", {})
        _post_json(
            f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat",
            {"completed_directives": [{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct"}]},
        )
        network = _get_json(f"http://127.0.0.1:{port}/v1/network/map")
        assert heartbeat["directives"][0]["action"] == "LOAD_SHARD"
        release = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/models/release",
            data=json.dumps({"model": "qwen2.5-0.5b-instruct"}).encode(), headers=headers, method="POST")
        with urllib.request.urlopen(release) as response:
            assert json.loads(response.read())["released"] == 1
        heartbeat = _post_json(f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat", {})
    finally:
        tracker.stop()

    assert heartbeat["directives"][0]["action"] == "DROP_SHARD"
    released_node = next(item for item in network["nodes"] if item["node_id"] == node["node_id"])
    assert released_node["shard_start"] is None
    assert released_node["shard_end"] is None


def test_models_list_does_not_duplicate_a_preset_registered_by_hf_repo():
    """A preset and its canonical repository are one selectable model."""
    tracker = TrackerServer(enable_billing=False)
    port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            {
                "endpoint": "http://127.0.0.1:9913",
                "model": "Qwen2.5-0.5B-Instruct",
                "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
                "num_layers": 24,
                "shard_start": 0,
                "shard_end": 23,
            },
        )
        models = _get_json(f"http://127.0.0.1:{port}/v1/models")["data"]
    finally:
        tracker.stop()

    assert [model["id"] for model in models].count("qwen2.5-0.5b-instruct") == 1
    assert not any(model["id"] == "Qwen/Qwen2.5-0.5B-Instruct" for model in models)


def test_endpoint_key_distinguishes_same_port_different_hosts():
    "Endpoint key distinguishes same port different hosts\n\nTags: http, performance, routing, tracker"
    from meshnet_node.torch_server import _clamp_downstream_hops, _endpoint_key

    assert _endpoint_key("http://192.168.0.20:7000") == "192.168.0.20:7000"
    assert _endpoint_key("http://192.168.0.179:7000") == "192.168.0.179:7000"
    assert _endpoint_key("http://192.168.0.20:7000") != _endpoint_key("http://192.168.0.179:7000")

    class Backend:
        shard_end = 21

    hops = [{"endpoint": "http://192.168.0.179:7000", "start_layer": 0}]
    assert _clamp_downstream_hops(hops, Backend()) == [
        {"endpoint": "http://192.168.0.179:7000", "start_layer": 22},
    ]
