"""Async HTTP client for mesh peer communication."""

from __future__ import annotations

import logging
import time
from typing import Any, Optional

import httpx

from radix_edge.mesh.security import sign_request

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


class PeerClient:
    """Async HTTP client for communicating with a single mesh peer."""

    def __init__(self, host: str, port: int, mesh_key: str):
        self.host = host
        self.port = port
        self.mesh_key = mesh_key
        self._client: Optional[httpx.AsyncClient] = None

    @property
    def base_url(self) -> str:
        return f"http://{self.host}:{self.port}"

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=10.0,
                headers={"Content-Type": "application/json"},
            )
        return self._client

    def _auth_headers(self, body: bytes) -> dict[str, str]:
        """Generate HMAC authentication headers."""
        timestamp = str(time.time())
        signature = sign_request(self.mesh_key, body, timestamp)
        return {
            "X-Radix-Mesh-Hmac": signature,
            "X-Radix-Mesh-Timestamp": timestamp,
        }

    async def send_heartbeat(self, node_id: str, capabilities: dict) -> dict[str, Any]:
        """Send heartbeat to this peer."""
        import json

        client = await self._get_client()
        payload = json.dumps({"node_id": node_id, "capabilities": capabilities}).encode("utf-8")
        resp = await client.post(
            "/v1/mesh/heartbeat",
            content=payload,
            headers=self._auth_headers(payload),
        )
        resp.raise_for_status()
        return resp.json()

    async def get_queued_jobs(self) -> list[dict]:
        """Get queued jobs from this peer."""
        client = await self._get_client()
        payload = b"{}"
        resp = await client.get(
            "/v1/mesh/peers",
            headers=self._auth_headers(payload),
        )
        # We actually need a dedicated endpoint for queued jobs;
        # use the jobs endpoint which returns job data
        resp = await client.post(
            "/v1/mesh/jobs/query",
            content=payload,
            headers=self._auth_headers(payload),
        )
        if resp.status_code == 404:
            # Peer may not support query endpoint; return empty
            return []
        resp.raise_for_status()
        return resp.json().get("jobs", [])

    async def assign_job(self, job_id: str, gpu_indices: list[int]) -> dict[str, Any]:
        """Assign a job to specific GPUs on this peer."""
        import json

        client = await self._get_client()
        payload = json.dumps({"job_id": job_id, "gpu_indices": gpu_indices}).encode("utf-8")
        resp = await client.post(
            "/v1/mesh/jobs/assign",
            content=payload,
            headers=self._auth_headers(payload),
        )
        resp.raise_for_status()
        return resp.json()

    async def forward_job(self, job_data: dict) -> dict[str, Any]:
        """Forward a job submission to this peer."""
        import json

        client = await self._get_client()
        payload = json.dumps(job_data).encode("utf-8")
        resp = await client.post(
            "/v1/mesh/jobs/forward",
            content=payload,
            headers=self._auth_headers(payload),
        )
        resp.raise_for_status()
        return resp.json()

    async def announce_election(self, coordinator_id: str) -> dict[str, Any]:
        """Announce election result to this peer."""
        import json

        client = await self._get_client()
        payload = json.dumps({"coordinator_id": coordinator_id}).encode("utf-8")
        resp = await client.post(
            "/v1/mesh/election/announce",
            content=payload,
            headers=self._auth_headers(payload),
        )
        resp.raise_for_status()
        return resp.json()

    async def close(self) -> None:
        """Close the underlying HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
