"""Async HTTP client for Radix control plane communication."""

from __future__ import annotations

import logging
from typing import Any, Optional

import httpx

logger = logging.getLogger("radix_edge.upstream")

AGENT_VERSION = "0.1.0"


class UpstreamClient:
    """Async HTTP client for the Radix cloud API."""

    def __init__(self, api_base: str, cluster_id: str, tenant_id: str, cluster_token: str):
        self.api_base = api_base.rstrip("/")
        self.cluster_id = cluster_id
        self.tenant_id = tenant_id
        self.cluster_token = cluster_token
        self._client: Optional[httpx.AsyncClient] = None

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.api_base,
                timeout=30.0,
                headers={
                    "X-Radix-Cluster-Token": self.cluster_token,
                    "X-Radix-Tenant-Id": self.tenant_id,
                    "Content-Type": "application/json",
                },
            )
        return self._client

    async def send_heartbeat(self, status: str, capabilities: dict[str, Any]) -> dict[str, Any]:
        """Send heartbeat to control plane."""
        client = await self._get_client()
        resp = await client.post(
            f"/v1/clusters/{self.cluster_id}/heartbeat",
            json={"status": status, "capabilities": capabilities, "agent_version": AGENT_VERSION},
        )
        resp.raise_for_status()
        return resp.json()

    async def poll_jobs(self, max_jobs: int = 1) -> list[dict[str, Any]]:
        """Poll for pending jobs from control plane."""
        client = await self._get_client()
        resp = await client.post(
            f"/v1/clusters/{self.cluster_id}/jobs/poll",
            json={"max_jobs": max_jobs},
        )
        resp.raise_for_status()
        return resp.json().get("jobs", [])

    async def complete_job(
        self,
        job_id: str,
        status: str,
        output_payload: Optional[dict[str, Any]] = None,
        metrics: Optional[dict[str, Any]] = None,
        error_message: Optional[str] = None,
    ) -> dict[str, Any]:
        """Report job completion to control plane."""
        client = await self._get_client()
        resp = await client.post(
            f"/v1/clusters/{self.cluster_id}/jobs/{job_id}/complete",
            json={
                "tenant_id": self.tenant_id,
                "status": status,
                "output_payload": output_payload or {},
                "metrics": metrics or {},
                "error_message": error_message,
            },
        )
        resp.raise_for_status()
        return resp.json()

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