"""US-030: manual route selection + hop-penalty benchmarking.

Pinned "route" in the chat body overrides auto-selection; the privileged
benchmark endpoint fans the same prompt through 1/2/3-node routes and appends
per-hop latency records to benchmark_results.json.
"""

import json
import urllib.error
import urllib.request

import pytest

from meshnet_node.server import StubNodeServer
from meshnet_tracker.server import TrackerServer

MODEL = "openai-community/gpt2"


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


@pytest.fixture
def benchmark_setup(tmp_path):
    """Tracker + one full-coverage node and one 0-5/6-11 pair, with node ids."""
    results_path = str(tmp_path / "benchmark_results.json")
    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        benchmark_results_path=results_path,
        validator_service_token="bench",
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"

    stubs = []
    node_ids = {}
    for name, (start, end, prefix, last) in {
        "full": (0, 11, "full:", True),
        "head": (0, 5, "head:", False),
        "tail": (6, 11, "tail:", True),
    }.items():
        stub = StubNodeServer(
            shard_start=start, shard_end=end, is_last_shard=last,
            response_prefix=prefix, model=MODEL, tracker_mode=(start == 0),
        )
        sport = stub.start()
        stubs.append(stub)
        reply = _post_json(f"{tracker_url}/v1/nodes/register", {
            "endpoint": f"http://127.0.0.1:{sport}",
            "shard_start": start,
            "shard_end": end,
            "model": MODEL,
            "hardware_profile": {},
            "score": 1.0,
            "tracker_mode": start == 0,
            "node_id": f"node-{name}",
        })
        node_ids[name] = reply.get("node_id", f"node-{name}")

    yield tracker_url, node_ids, results_path

    for stub in stubs:
        stub.stop()
    tracker.stop()


def _chat(tracker_url: str, route: list[str] | None = None) -> dict:
    body: dict = {
        "model": MODEL,
        "messages": [{"role": "user", "content": "ping"}],
    }
    if route is not None:
        body["route"] = route
    return _post_json(f"{tracker_url}/v1/chat/completions", body)


def test_pinned_route_uses_named_node(benchmark_setup):
    "Pinned route uses named node\n\nTags: routing"
    tracker_url, node_ids, _ = benchmark_setup
    reply = _chat(tracker_url, route=[node_ids["full"]])
    content = reply["choices"][0]["message"]["content"]
    assert content.startswith("full:")


def test_unknown_route_node_is_400(benchmark_setup):
    "Unknown route node is 400\n\nTags: routing"
    tracker_url, _, _ = benchmark_setup
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _chat(tracker_url, route=["no-such-node"])
    assert exc_info.value.code == 400
    detail = json.loads(exc_info.value.read())
    assert "no-such-node" in detail["error"]["message"]


def test_invalid_route_shape_is_400(benchmark_setup):
    "Invalid route shape is 400\n\nTags: routing"
    tracker_url, _, _ = benchmark_setup
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _chat(tracker_url, route=[])
    assert exc_info.value.code == 400


def test_clients_without_route_are_unaffected(benchmark_setup):
    "Clients without route are unaffected\n\nTags: routing"
    tracker_url, _, _ = benchmark_setup
    reply = _chat(tracker_url)
    assert reply["choices"][0]["message"]["content"]


def test_benchmark_requires_auth(benchmark_setup):
    "Benchmark requires auth\n\nTags: auth, performance, security"
    tracker_url, _, _ = benchmark_setup
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _post_json(f"{tracker_url}/v1/benchmark/hop-penalty", {"model": MODEL})
    assert exc_info.value.code == 401
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        urllib.request.urlopen(f"{tracker_url}/v1/benchmark/results")
    assert exc_info.value.code == 401


def test_benchmark_records_one_and_two_node_routes(benchmark_setup):
    "Benchmark records one and two node routes\n\nTags: performance, routing"
    tracker_url, _, results_path = benchmark_setup
    record = _post_json(
        f"{tracker_url}/v1/benchmark/hop-penalty",
        {"model": MODEL, "prompt": "2+2", "max_new_tokens": 8},
        headers={"Authorization": "Bearer bench"},
    )
    by_hops = {len(r["route"]): r for r in record["routes"] if "total_ms" in r}
    assert 1 in by_hops and 2 in by_hops  # 3-node coverage impossible → skipped
    assert 3 not in by_hops
    assert len(by_hops[1]["per_hop_ms"]) == 1
    assert len(by_hops[2]["per_hop_ms"]) == 2
    assert by_hops[2]["total_ms"] > 0

    stored = json.loads(open(results_path, encoding="utf-8").read())
    assert isinstance(stored, list) and len(stored) == 1
    assert stored[0]["model"] == MODEL
    assert stored[0]["prompt_hash"]

    fetched = json.loads(
        urllib.request.urlopen(urllib.request.Request(
            f"{tracker_url}/v1/benchmark/results",
            headers={"Authorization": "Bearer bench"},
        )).read()
    )
    assert len(fetched["results"]) == 1