"""HTTPS client for the Specter broker API.

Handles provisioning, session status, teardown, and heartbeat
calls to ``broker.specter.host``.
"""

from __future__ import annotations

import logging
import os
from typing import Any

import httpx

log = logging.getLogger(__name__)

DEFAULT_BROKER_URL = "https://broker.specter.host"


def get_broker_url() -> str:
    """Return the broker URL, reading ``SPECTER_BROKER_URL`` env var if set."""
    return os.environ.get("SPECTER_BROKER_URL", DEFAULT_BROKER_URL)


class BrokerError(Exception):
    """Raised when the broker returns an error response."""

    def __init__(self, status_code: int, error: str, message: str) -> None:
        self.status_code = status_code
        self.error = error
        super().__init__(f"Broker error ({status_code}): {error} — {message}")


class BrokerClient:
    """Client for the Specter broker provisioning API.

    Parameters
    ----------
    base_url:
        Broker API base URL. Defaults to ``https://broker.specter.host``.
    token:
        Authentication token from ``specter login``.
    timeout:
        Request timeout in seconds.
    """

    def __init__(
        self,
        token: str,
        *,
        base_url: str = DEFAULT_BROKER_URL,
        timeout: float = 30.0,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._client = httpx.Client(
            base_url=f"{self._base_url}/v1",
            headers={"Authorization": f"Bearer {token}"},
            timeout=timeout,
        )

    def _handle_response(self, response: httpx.Response) -> dict[str, Any]:
        """Parse response and raise BrokerError on failure.

        FastAPI wraps HTTPException errors in ``{"detail": {...}}``, so we
        check for that wrapper first before reading top-level keys.
        """
        if response.status_code >= 400:
            try:
                body = response.json()
            except Exception:
                body = {"error": "unknown", "message": response.text}
            # FastAPI wraps HTTPException detail in {"detail": {...}}.
            detail = body.get("detail", body) if isinstance(body, dict) else body
            if isinstance(detail, dict):
                error = detail.get("error", "unknown")
                message = detail.get("message", "")
            else:
                error = "unknown"
                message = str(detail)
            raise BrokerError(response.status_code, error, message)
        return response.json()

    def provision(
        self,
        gpu_type: str,
        ssh_public_key: str,
        *,
        gpu_count: int = 1,
        disk_gb: int = 20,
        image: str | None = None,
    ) -> dict[str, Any]:
        """Provision a new GPU session.

        Returns session info with ``session_id`` and ``status``.
        """
        payload: dict[str, Any] = {
            "gpu_type": gpu_type,
            "ssh_public_key": ssh_public_key,
            "gpu_count": gpu_count,
            "disk_gb": disk_gb,
        }
        if image:
            payload["image"] = image

        log.debug("Provisioning %s GPU via broker", gpu_type)
        resp = self._client.post("/provision", json=payload)
        return self._handle_response(resp)

    def get_session(self, session_id: str) -> dict[str, Any]:
        """Get session status and connection details."""
        resp = self._client.get(f"/sessions/{session_id}")
        return self._handle_response(resp)

    def destroy_session(self, session_id: str) -> dict[str, Any]:
        """Tear down a session. Idempotent."""
        resp = self._client.delete(f"/sessions/{session_id}")
        return self._handle_response(resp)

    def heartbeat(self, session_id: str) -> dict[str, Any]:
        """Send keepalive heartbeat for a session."""
        resp = self._client.post("/heartbeat", json={"session_id": session_id})
        return self._handle_response(resp)

    def list_gpu_types(self) -> list[dict[str, Any]]:
        """Fetch available GPU types from the broker.

        Calls ``GET /v1/gpu-types`` (public endpoint, no auth needed).
        Returns a list of GPU type objects with at least a ``name`` field.
        """
        resp = self._client.get("/gpu-types")
        data = self._handle_response(resp)
        return data.get("gpu_types", [])

    def close(self) -> None:
        """Close the underlying HTTP client."""
        self._client.close()

    def __enter__(self) -> BrokerClient:
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()
