"""AH-021: tracker-scheduled TOPLOC honest-noise calibration dispatch.

Extends the US-030 fleet-dispatch pattern (`_handle_benchmark_hop_penalty`)
from pinned-route latency benchmarking to a job that hits every solo-capable
registered node with a fixed prompt, verifies each node's own on-demand
TOPLOC commitment against a teacher-forced reference replay, and records the
raw divergence into the calibration corpus.
"""

from __future__ import annotations

import http.server
import json
import socketserver
import threading
import time
import urllib.error
import urllib.request

import pytest

from meshnet_tracker.server import TrackerServer
from meshnet_validator.audit import ToplocAuditConfig, build_activation_proofs

MODEL = "openai-community/gpt2"
CONFIG = ToplocAuditConfig(topk=2, decode_batching_size=16, dtype="bfloat16", quantization="bfloat16")


class FakeToploc:
    """Exact-equality fake backend, matching other TOPLOC test suites."""

    def build_proofs_base64(self, activations, *, decode_batching_size, topk, skip_prefill):
        return {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }

    def verify_proofs_base64(self, activations, proofs, *, decode_batching_size, topk, skip_prefill):
        # `proofs` may have round-tripped through JSON (HTTP dispatch), which
        # turns the fingerprint's tuples into lists — normalize before compare.
        fingerprint = proofs.get("activation_fingerprint") if isinstance(proofs, dict) else None
        return (
            fingerprint is not None
            and tuple(tuple(row) for row in fingerprint) == tuple(tuple(row) for row in activations)
            and proofs.get("decode_batching_size") == decode_batching_size
            and proofs.get("topk") == topk
            and proofs.get("skip_prefill") == skip_prefill
        )


BACKEND = FakeToploc()


class FakeCalibrationNode:
    """Stands in for a node: serves both /v1/chat/completions (tracker_mode
    style) and its own on-demand TOPLOC commitment endpoint."""

    def __init__(self, *, claim_activations, claimed_token_ids, response_text="ok", commitment_available=True):
        self.requests: list[dict] = []
        claim = build_activation_proofs(claim_activations, config=CONFIG, backend=BACKEND)
        outer = self

        class Handler(http.server.BaseHTTPRequestHandler):
            def log_message(self, fmt, *args):
                pass

            def do_POST(self):
                length = int(self.headers.get("Content-Length", 0))
                body = json.loads(self.rfile.read(length) or b"{}")
                if self.path == "/v1/chat/completions":
                    self._send_json(200, {
                        "id": "chatcmpl-cal",
                        "object": "chat.completion",
                        "created": int(time.time()),
                        "model": body.get("model", MODEL),
                        "choices": [{
                            "index": 0,
                            "message": {"role": "assistant", "content": response_text},
                            "finish_reason": "stop",
                        }],
                        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
                    })
                    return
                if self.path == "/v1/audit/toploc/commitment":
                    outer.requests.append(body)
                    if not commitment_available:
                        self.send_response(404)
                        self.end_headers()
                        return
                    self._send_json(200, {
                        "toploc_proof": claim.as_mapping(),
                        "claimed_token_ids": claimed_token_ids,
                    })
                    return
                self.send_response(404)
                self.end_headers()

            def _send_json(self, status, data):
                payload = json.dumps(data).encode()
                self.send_response(status)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)

        self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler)
        self._server.daemon_threads = True
        self._thread: threading.Thread | None = None

    def start(self) -> str:
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return f"http://127.0.0.1:{self._server.server_address[1]}"

    def stop(self) -> None:
        self._server.shutdown()
        self._server.server_close()


class FakeReferenceNode:
    """Stands in for the reference node: teacher-forces the claimed tokens
    and returns canned reference activations."""

    def __init__(self, *, reference_activations):
        self.requests: list[dict] = []
        outer = self

        class Handler(http.server.BaseHTTPRequestHandler):
            def log_message(self, fmt, *args):
                pass

            def do_POST(self):
                if self.path != "/v1/audit/toploc":
                    self.send_response(404)
                    self.end_headers()
                    return
                length = int(self.headers.get("Content-Length", 0))
                body = json.loads(self.rfile.read(length) or b"{}")
                outer.requests.append(body)
                payload = json.dumps({"activations": reference_activations}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)

        self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler)
        self._server.daemon_threads = True
        self._thread: threading.Thread | None = None

    def start(self) -> str:
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return f"http://127.0.0.1:{self._server.server_address[1]}"

    def stop(self) -> None:
        self._server.shutdown()
        self._server.server_close()


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())


def _get_json(url: str, headers: dict | None = None) -> dict:
    req = urllib.request.Request(url, headers=headers or {})
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


@pytest.fixture
def calibration_setup(tmp_path):
    reference_activations = [[1.0, 2.0], [3.0, 4.0]]
    reference = FakeReferenceNode(reference_activations=reference_activations)
    reference_url = reference.start()

    calibration_db = str(tmp_path / "calibration.sqlite")
    tracker = TrackerServer(
        model_presets={
            MODEL: {"layers_start": 0, "layers_end": 11, "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}},
        },
        validator_service_token="cal-token",
        toploc_calibration_db=calibration_db,
        toploc_reference_node_url=reference_url,
        toploc_calibration_gate_min_hardware_profiles=1,
        toploc_backend=BACKEND,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"

    honest_node = FakeCalibrationNode(
        claim_activations=reference_activations,  # matches reference -> passes
        claimed_token_ids=[101, 202],
    )
    honest_node_url = honest_node.start()
    reply = _post_json(f"{tracker_url}/v1/nodes/register", {
        "endpoint": honest_node_url,
        "shard_start": 0,
        "shard_end": 11,
        "model": MODEL,
        "hardware_profile": {"gpu_name": "RTX 4090"},
        "quantization": "bfloat16",
        "wallet_address": "wallet-honest",
        "score": 1.0,
        "tracker_mode": True,
    })
    honest_node_id = reply["node_id"]

    partial_node = FakeCalibrationNode(claim_activations=reference_activations, claimed_token_ids=[101])
    partial_node_url = partial_node.start()
    reply = _post_json(f"{tracker_url}/v1/nodes/register", {
        "endpoint": partial_node_url,
        "shard_start": 0,
        "shard_end": 5,
        "model": MODEL,
        "hardware_profile": {"gpu_name": "A100"},
        "quantization": "bfloat16",
        "wallet_address": "wallet-partial",
        "score": 1.0,
        "tracker_mode": True,
    })
    partial_node_id = reply["node_id"]

    yield tracker_url, calibration_db, honest_node_id, partial_node_id

    honest_node.stop()
    partial_node.stop()
    reference.stop()
    tracker.stop()


def test_calibration_run_requires_auth(calibration_setup):
    "Calibration run requires auth\n\nTags: audit, auth, calibration, security"
    tracker_url, _, _, _ = calibration_setup
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _post_json(f"{tracker_url}/v1/calibration/toploc/run", {"model": MODEL})
    assert exc_info.value.code == 401


def test_calibration_run_dispatches_only_solo_capable_nodes(calibration_setup):
    "Calibration run dispatches only solo capable nodes\n\nTags: audit, calibration"
    tracker_url, _, honest_node_id, partial_node_id = calibration_setup
    record = _post_json(
        f"{tracker_url}/v1/calibration/toploc/run",
        {"model": MODEL, "prompt": "2+2", "max_new_tokens": 4},
        headers={"Authorization": "Bearer cal-token"},
    )
    assert record["skipped_partial_shard_node_ids"] == [partial_node_id]
    assert len(record["nodes"]) == 1
    result = record["nodes"][0]
    assert result["node_id"] == honest_node_id
    assert result["wallet_address"] == "wallet-honest"
    assert result["gpu_model"] == "RTX 4090"
    assert result["dtype"] == "bfloat16"
    assert result["passed"] is True


def test_calibration_run_persists_corpus_and_results_endpoint_reports_it(calibration_setup):
    "Calibration run persists corpus and results endpoint reports it\n\nTags: audit, calibration, persistence"
    tracker_url, calibration_db, _, _ = calibration_setup
    _post_json(
        f"{tracker_url}/v1/calibration/toploc/run",
        {"model": MODEL},
        headers={"Authorization": "Bearer cal-token"},
    )
    results = _get_json(
        f"{tracker_url}/v1/calibration/toploc/results",
        headers={"Authorization": "Bearer cal-token"},
    )
    assert len(results["runs"]) == 1
    assert results["runs"][0]["node_wallet"] == "wallet-honest"
    assert results["gate_status"]["distinct_hardware_profiles"] == 1
    assert results["gate_status"]["ready"] is True
    assert results["envelope"]["sample_count"] == 1


def test_calibration_run_missing_reference_node_url_is_503(tmp_path):
    "Calibration run missing reference node url is 503\n\nTags: audit, calibration"
    tracker = TrackerServer(
        model_presets={MODEL: {"layers_start": 0, "layers_end": 11}},
        validator_service_token="cal-token",
        toploc_calibration_db=str(tmp_path / "calibration.sqlite"),
    )
    port = tracker.start()
    try:
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            _post_json(
                f"http://127.0.0.1:{port}/v1/calibration/toploc/run",
                {"model": MODEL},
                headers={"Authorization": "Bearer cal-token"},
            )
        assert exc_info.value.code == 503
    finally:
        tracker.stop()


def test_calibration_run_node_without_commitment_endpoint_is_skipped_not_failed(tmp_path):
    "Calibration run node without commitment endpoint is skipped not failed\n\nTags: audit, calibration"
    reference = FakeReferenceNode(reference_activations=[[1.0, 2.0]])
    reference_url = reference.start()
    tracker = TrackerServer(
        model_presets={MODEL: {"layers_start": 0, "layers_end": 11}},
        validator_service_token="cal-token",
        toploc_calibration_db=str(tmp_path / "calibration.sqlite"),
        toploc_reference_node_url=reference_url,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"
    node = FakeCalibrationNode(
        claim_activations=[[1.0, 2.0]], claimed_token_ids=[101], commitment_available=False,
    )
    node_url = node.start()
    _post_json(f"{tracker_url}/v1/nodes/register", {
        "endpoint": node_url,
        "shard_start": 0,
        "shard_end": 11,
        "model": MODEL,
        "hardware_profile": {},
        "wallet_address": "wallet-no-commitment",
        "score": 1.0,
        "tracker_mode": True,
        "node_id": "node-no-commitment",
    })

    try:
        record = _post_json(
            f"{tracker_url}/v1/calibration/toploc/run",
            {"model": MODEL},
            headers={"Authorization": "Bearer cal-token"},
        )
        assert len(record["nodes"]) == 1
        assert "skipped" in record["nodes"][0]
        assert record["gate_status"]["sample_count"] == 0
    finally:
        node.stop()
        reference.stop()
        tracker.stop()