"""FastAPI router for mesh peer-to-peer endpoints."""

from __future__ import annotations

import json
import logging
from typing import Any

from fastapi import APIRouter, HTTPException, Request

from radix_edge.mesh.security import verify_request

logger = logging.getLogger("radix_edge.mesh.routes")

router = APIRouter(prefix="/v1/mesh", tags=["mesh"])


def _get_mesh_state(request: Request) -> dict[str, Any]:
    """Retrieve mesh objects from app.state."""
    return {
        "registry": getattr(request.app.state, "mesh_registry", None),
        "coordinator": getattr(request.app.state, "mesh_coordinator", None),
        "mesh_key": getattr(request.app.state, "mesh_key", ""),
        "node_id": getattr(request.app.state, "mesh_node_id", ""),
    }


async def _validate_hmac(request: Request, mesh_key: str) -> bytes:
    """Validate HMAC signature on incoming request. Returns the raw body."""
    body = await request.body()
    signature = request.headers.get("X-Radix-Mesh-Hmac", "")
    timestamp = request.headers.get("X-Radix-Mesh-Timestamp", "")

    if not signature or not timestamp:
        raise HTTPException(status_code=401, detail="Missing mesh authentication headers")

    if not verify_request(mesh_key, body, timestamp, signature):
        raise HTTPException(status_code=403, detail="Invalid mesh signature or expired timestamp")

    return body


@router.post("/heartbeat")
async def receive_heartbeat(request: Request) -> dict[str, Any]:
    """Receive a heartbeat from a mesh peer."""
    state = _get_mesh_state(request)
    body = await _validate_hmac(request, state["mesh_key"])
    data = json.loads(body)

    node_id = data.get("node_id", "")
    capabilities = data.get("capabilities", {})

    if not node_id:
        raise HTTPException(status_code=400, detail="Missing node_id in heartbeat")

    registry = state["registry"]
    if registry is None:
        raise HTTPException(status_code=503, detail="Mesh not initialized")

    # Determine peer host from the request
    host = request.client.host if request.client else "unknown"
    port = capabilities.get("port", 8080)

    registry.update_peer(node_id, host, port, capabilities)
    logger.debug("Heartbeat received from %s", node_id)

    return {"status": "ok", "coordinator_id": state["coordinator"].coordinator_id if state["coordinator"] else None}


@router.post("/jobs/assign")
async def receive_job_assignment(request: Request) -> dict[str, Any]:
    """Receive a job assignment from the coordinator."""
    state = _get_mesh_state(request)
    body = await _validate_hmac(request, state["mesh_key"])
    data = json.loads(body)

    job_id = data.get("job_id", "")
    gpu_indices = data.get("gpu_indices", [])

    if not job_id:
        raise HTTPException(status_code=400, detail="Missing job_id")

    # Update the job's target_node to this node and mark for local scheduling
    from radix_edge.db import get_db

    with get_db() as db:
        row = db.execute("SELECT status FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")

        db.execute(
            "UPDATE jobs SET target_node = ? WHERE job_id = ?",
            (state["node_id"], job_id),
        )

    logger.info("Job %s assigned to this node with GPUs %s", job_id, gpu_indices)
    return {"status": "accepted", "job_id": job_id}


@router.post("/jobs/forward")
async def forward_job(request: Request) -> dict[str, Any]:
    """Receive a forwarded job from another peer."""
    state = _get_mesh_state(request)
    body = await _validate_hmac(request, state["mesh_key"])
    data = json.loads(body)

    # Insert the job into local queue
    import uuid

    from radix_edge.db import get_db

    job_id = data.get("job_id") or f"job-{uuid.uuid4().hex[:12]}"
    image = data.get("image", "")
    user_id = data.get("user_id", "mesh")

    if not image:
        raise HTTPException(status_code=400, detail="Missing image in forwarded job")

    with get_db() as db:
        db.execute(
            """INSERT OR IGNORE INTO jobs (job_id, user_id, image, job_type, priority, num_gpus, origin, target_node)
               VALUES (?, ?, ?, ?, ?, ?, 'mesh', ?)""",
            (
                job_id,
                user_id,
                image,
                data.get("job_type", "generic"),
                data.get("priority", 5),
                data.get("num_gpus", 1),
                state["node_id"],
            ),
        )

    logger.info("Forwarded job %s inserted into local queue", job_id)
    return {"status": "accepted", "job_id": job_id}


@router.post("/jobs/query")
async def query_queued_jobs(request: Request) -> dict[str, Any]:
    """Return queued jobs on this node (used by coordinator for scheduling)."""
    state = _get_mesh_state(request)
    await _validate_hmac(request, state["mesh_key"])

    from radix_edge.db import get_db

    with get_db() as db:
        rows = db.execute(
            "SELECT job_id, user_id, image, job_type, priority, num_gpus, gpu_mem_required_mb "
            "FROM jobs WHERE status = 'queued' ORDER BY created_at LIMIT 50"
        ).fetchall()
        jobs = [dict(r) for r in rows]

    return {"jobs": jobs}


@router.post("/election/announce")
async def election_announce(request: Request) -> dict[str, Any]:
    """Receive an election result announcement."""
    state = _get_mesh_state(request)
    body = await _validate_hmac(request, state["mesh_key"])
    data = json.loads(body)

    coordinator_id = data.get("coordinator_id", "")
    if not coordinator_id:
        raise HTTPException(status_code=400, detail="Missing coordinator_id")

    coordinator = state["coordinator"]
    if coordinator is None:
        raise HTTPException(status_code=503, detail="Mesh not initialized")

    coordinator.accept_coordinator(coordinator_id)
    logger.info("Election announcement: coordinator is %s", coordinator_id)

    return {"status": "ok", "accepted_coordinator": coordinator_id}


@router.get("/peers")
async def list_peers(request: Request) -> dict[str, Any]:
    """List all known mesh peers."""
    state = _get_mesh_state(request)
    registry = state["registry"]
    if registry is None:
        return {"peers": [], "node_id": state["node_id"]}

    peers = []
    for p in registry.get_active_peers():
        peers.append({
            "node_id": p.node_id,
            "host": p.host,
            "port": p.port,
            "capabilities": p.capabilities,
            "is_coordinator": p.is_coordinator,
        })

    return {"peers": peers, "node_id": state["node_id"]}


@router.get("/status")
async def mesh_status(request: Request) -> dict[str, Any]:
    """Return mesh status summary."""
    state = _get_mesh_state(request)
    registry = state["registry"]
    coordinator = state["coordinator"]

    active_count = len(registry.get_active_peers()) if registry else 0
    stale_count = len(registry.get_stale_peers()) if registry else 0

    return {
        "node_id": state["node_id"],
        "mesh_enabled": True,
        "coordinator_id": coordinator.coordinator_id if coordinator else None,
        "is_coordinator": coordinator.is_coordinator if coordinator else False,
        "active_peers": active_count,
        "stale_peers": stale_count,
        "total_gpus": len(registry.get_aggregate_gpus()) if registry else 0,
    }
