"""Tests for mesh API routes."""

from __future__ import annotations

import json
import time

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from radix_edge.mesh.coordinator import MeshCoordinator
from radix_edge.mesh.peer_registry import PeerRegistry
from radix_edge.mesh.routes_mesh import router
from radix_edge.mesh.security import sign_request


def _create_mesh_app(db_path: str) -> FastAPI:
    """Create a minimal FastAPI app with mesh routes for testing."""
    app = FastAPI()
    app.include_router(router)

    # Set up mesh state
    registry = PeerRegistry()
    coordinator = MeshCoordinator("node-test", registry)
    coordinator.elect()

    app.state.mesh_registry = registry
    app.state.mesh_coordinator = coordinator
    app.state.mesh_key = "test-mesh-key"
    app.state.mesh_node_id = "node-test"

    return app


def _auth_headers(body: bytes, mesh_key: str = "test-mesh-key") -> dict:
    """Generate valid HMAC headers for testing."""
    ts = str(time.time())
    sig = sign_request(mesh_key, body, ts)
    return {
        "X-Radix-Mesh-Hmac": sig,
        "X-Radix-Mesh-Timestamp": ts,
    }


class TestHeartbeatEndpoint:
    def test_receive_heartbeat(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"node_id": "node-peer", "capabilities": {"gpu_count": 2}}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/heartbeat", content=body, headers=_auth_headers(body))
            assert resp.status_code == 200
            data = resp.json()
            assert data["status"] == "ok"

    def test_heartbeat_missing_node_id(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"capabilities": {}}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/heartbeat", content=body, headers=_auth_headers(body))
            assert resp.status_code == 400

    def test_heartbeat_invalid_hmac(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"node_id": "node-peer"}).encode()

        with TestClient(app) as client:
            resp = client.post(
                "/v1/mesh/heartbeat",
                content=body,
                headers={"X-Radix-Mesh-Hmac": "bad", "X-Radix-Mesh-Timestamp": str(time.time())},
            )
            assert resp.status_code == 403

    def test_heartbeat_missing_hmac_headers(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"node_id": "node-peer"}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/heartbeat", content=body)
            assert resp.status_code == 401


class TestJobAssignEndpoint:
    def test_assign_job(self, db_path):
        app = _create_mesh_app(db_path)

        # Insert a job first
        from radix_edge.db import get_db
        with get_db(db_path) as db:
            db.execute(
                "INSERT INTO jobs (job_id, user_id, image) VALUES (?, ?, ?)",
                ("job-mesh-1", "user-1", "img:1"),
            )

        body = json.dumps({"job_id": "job-mesh-1", "gpu_indices": [0, 1]}).encode()
        with TestClient(app) as client:
            resp = client.post("/v1/mesh/jobs/assign", content=body, headers=_auth_headers(body))
            assert resp.status_code == 200
            assert resp.json()["status"] == "accepted"

    def test_assign_nonexistent_job(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"job_id": "job-ghost", "gpu_indices": [0]}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/jobs/assign", content=body, headers=_auth_headers(body))
            assert resp.status_code == 404


class TestForwardJobEndpoint:
    def test_forward_job(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({
            "job_id": "job-fwd-1",
            "user_id": "user-remote",
            "image": "nvidia/cuda:12.4",
            "num_gpus": 1,
        }).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/jobs/forward", content=body, headers=_auth_headers(body))
            assert resp.status_code == 200
            assert resp.json()["job_id"] == "job-fwd-1"

    def test_forward_job_missing_image(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"user_id": "user-1"}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/jobs/forward", content=body, headers=_auth_headers(body))
            assert resp.status_code == 400


class TestQueryEndpoint:
    def test_query_queued_jobs(self, db_path):
        app = _create_mesh_app(db_path)

        from radix_edge.db import get_db
        with get_db(db_path) as db:
            db.execute(
                "INSERT INTO jobs (job_id, user_id, image, status) VALUES (?, ?, ?, ?)",
                ("job-q-1", "user-1", "img:1", "queued"),
            )
            db.execute(
                "INSERT INTO jobs (job_id, user_id, image, status) VALUES (?, ?, ?, ?)",
                ("job-r-1", "user-1", "img:1", "running"),
            )

        body = b"{}"
        with TestClient(app) as client:
            resp = client.post("/v1/mesh/jobs/query", content=body, headers=_auth_headers(body))
            assert resp.status_code == 200
            jobs = resp.json()["jobs"]
            assert len(jobs) == 1
            assert jobs[0]["job_id"] == "job-q-1"


class TestElectionAnnounceEndpoint:
    def test_election_announce(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({"coordinator_id": "node-new-coord"}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/election/announce", content=body, headers=_auth_headers(body))
            assert resp.status_code == 200
            assert resp.json()["accepted_coordinator"] == "node-new-coord"

    def test_election_announce_missing_id(self, db_path):
        app = _create_mesh_app(db_path)
        body = json.dumps({}).encode()

        with TestClient(app) as client:
            resp = client.post("/v1/mesh/election/announce", content=body, headers=_auth_headers(body))
            assert resp.status_code == 400


class TestPeersEndpoint:
    def test_list_peers_empty(self, db_path):
        app = _create_mesh_app(db_path)
        with TestClient(app) as client:
            resp = client.get("/v1/mesh/peers")
            assert resp.status_code == 200
            assert resp.json()["peers"] == []
            assert resp.json()["node_id"] == "node-test"

    def test_list_peers_with_data(self, db_path):
        app = _create_mesh_app(db_path)
        app.state.mesh_registry.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 2})

        with TestClient(app) as client:
            resp = client.get("/v1/mesh/peers")
            assert resp.status_code == 200
            peers = resp.json()["peers"]
            assert len(peers) == 1
            assert peers[0]["node_id"] == "node-a"


class TestStatusEndpoint:
    def test_mesh_status(self, db_path):
        app = _create_mesh_app(db_path)
        with TestClient(app) as client:
            resp = client.get("/v1/mesh/status")
            assert resp.status_code == 200
            data = resp.json()
            assert data["mesh_enabled"] is True
            assert data["node_id"] == "node-test"
            assert data["is_coordinator"] is True
