"""HTTP client for the Velar API."""

from __future__ import annotations

import time
from typing import Any

import httpx

from velar.config import Config

# Terminal deployment states that will never transition to "running".
_TERMINAL_STATES = frozenset({"completed", "failed", "cancelled", "error"})


class RemoteExecutionError(Exception):
    """Raised when the remote function raises an exception on the GPU."""


class DeploymentPollError(Exception):
    """Raised when a deployment reaches a terminal state other than 'running'."""


class DeploymentPollTimeout(Exception):
    """Raised when polling a deployment exceeds the timeout."""


class VelarClient:
    """Low-level HTTP client for the Velar API."""

    def __init__(self, config: Config | None = None):
        self._config = config or Config.from_env()
        self._client = httpx.Client(
            base_url=self._config.api_url,
            headers={
                "X-API-Key": self._config.api_key,
                "Content-Type": "application/json",
            },
            timeout=30.0,
        )

    # ------------------------------------------------------------------
    # Deployment CRUD
    # ------------------------------------------------------------------

    def create_deployment(
        self,
        *,
        deployment_type: str,
        gpu_type: str,
        image_url: str,
        entry_command: str = "",
        estimated_seconds: int = 0,
        env_vars: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        """Create a new deployment."""
        body: dict[str, Any] = {
            "type": deployment_type,
            "gpu_type": gpu_type,
            "image_url": image_url,
            "entry_command": entry_command,
        }
        if estimated_seconds > 0:
            body["estimated_seconds"] = estimated_seconds
        if env_vars:
            body["env_vars"] = env_vars

        resp = self._client.post("/api/v1/deployments", json=body)
        resp.raise_for_status()
        return resp.json()

    def get_deployment(self, deployment_id: str) -> dict[str, Any]:
        """Get deployment status."""
        resp = self._client.get(f"/api/v1/deployments/{deployment_id}")
        resp.raise_for_status()
        return resp.json()

    def cancel_deployment(self, deployment_id: str) -> dict[str, Any]:
        """Cancel a deployment."""
        resp = self._client.delete(f"/api/v1/deployments/{deployment_id}")
        resp.raise_for_status()
        return resp.json()

    def list_deployments(self, limit: int = 20, offset: int = 0) -> list[dict[str, Any]]:
        """List user deployments."""
        resp = self._client.get(
            "/api/v1/deployments", params={"limit": limit, "offset": offset}
        )
        resp.raise_for_status()
        return resp.json().get("deployments", [])

    # ------------------------------------------------------------------
    # Remote invocation helpers
    # ------------------------------------------------------------------

    def poll_deployment(
        self,
        deployment_id: str,
        *,
        timeout: float = 300,
        interval: float = 2,
    ) -> dict[str, Any]:
        """Poll a deployment until it reaches the 'running' state.

        Args:
            deployment_id: The deployment to poll.
            timeout: Maximum seconds to wait before raising.
            interval: Seconds between poll requests.

        Returns:
            The deployment dict once it is running.

        Raises:
            DeploymentPollTimeout: If the timeout is exceeded.
            DeploymentPollError: If the deployment reaches a terminal failure state.
        """
        deadline = time.monotonic() + timeout

        while True:
            deployment = self.get_deployment(deployment_id)
            status = deployment.get("status", "")

            if status == "running":
                return deployment

            if status in _TERMINAL_STATES:
                raise DeploymentPollError(
                    f"Deployment {deployment_id} reached terminal state '{status}' "
                    f"instead of 'running'."
                )

            if time.monotonic() >= deadline:
                raise DeploymentPollTimeout(
                    f"Deployment {deployment_id} did not reach 'running' within "
                    f"{timeout}s (last status: '{status}')."
                )

            time.sleep(interval)

    def invoke_function(
        self,
        deployment_id: str,
        payload: str,
        *,
        timeout: float = 600,
    ) -> str:
        """POST a serialized payload to a running deployment via invoke.velar.run.

        Args:
            deployment_id: The deployment ID to invoke.
            payload: A JSON string produced by :func:`velar.serialization.serialize`.
            timeout: HTTP request timeout in seconds.

        Returns:
            The raw response body as a string.

        Raises:
            RemoteExecutionError: If the remote function raises an exception.
        """
        invoke_url = f"{self._config.invoke_url}/{deployment_id}"
        resp = httpx.post(
            invoke_url,
            content=payload,
            headers={
                "Content-Type": "application/json",
                "X-API-Key": self._config.api_key,
            },
            timeout=timeout,
        )
        # Parse error from handler before raise_for_status so the user sees
        # the actual Python traceback from the remote function.
        if resp.status_code == 500:
            try:
                body = resp.json()
                if "error" in body:
                    # Print any stdout captured before the exception.
                    for line in body.get("logs") or []:
                        print(f"[velar] {line}")
                    raise RemoteExecutionError(
                        f"Remote function raised an exception:\n{body['error']}"
                    )
            except (ValueError, KeyError):
                pass
        resp.raise_for_status()
        return resp.text

    def get_deployment_logs(self, deployment_id: str) -> dict[str, Any]:
        """Get logs and live pod status for a deployment."""
        resp = self._client.get(f"/api/v1/deployments/{deployment_id}/logs")
        resp.raise_for_status()
        return resp.json()

    # ------------------------------------------------------------------
    # GPU pricing
    # ------------------------------------------------------------------

    def get_gpu_prices(self) -> list[dict[str, Any]]:
        """Fetch live GPU pricing from the Velar API.

        Returns a list of dicts with keys: ``gpu_type``, ``display_name``,
        ``vram_gb``, ``price_per_hour``, ``price_per_second``, ``provider``.
        """
        resp = self._client.get("/api/v1/gpu/prices")
        resp.raise_for_status()
        return resp.json().get("prices", [])

    # ------------------------------------------------------------------
    # User / billing
    # ------------------------------------------------------------------

    def get_me(self) -> dict[str, Any]:
        """Get current user info."""
        resp = self._client.get("/api/v1/users/me")
        resp.raise_for_status()
        return resp.json()

    def get_balance(self) -> dict[str, Any]:
        """Get credit balance."""
        resp = self._client.get("/api/v1/billing/balance")
        resp.raise_for_status()
        return resp.json()

    def close(self) -> None:
        self._client.close()
