"""RunPod serverless backend: submit a job, read its status."""

from __future__ import annotations

import httpx

from .sim_backend import BackendError


class RunPodError(BackendError):
    pass


class RunPodClient:
    name = "runpod"
    FUNCTIONS = {"run_simulation", "sweep"}  # what worker/Dockerfile's handler can run

    def __init__(
        self, api_key: str, endpoint_id: str, transport: httpx.BaseTransport | None = None, timeout: float = 30.0
    ):
        if not api_key or not endpoint_id:
            raise RunPodError("RUNPOD_API_KEY and RUNPOD_ENDPOINT_ID must be set")
        self._client = httpx.Client(
            base_url=f"https://api.runpod.ai/v2/{endpoint_id}",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=timeout,
            transport=transport,
        )

    def supports(self, function: str) -> bool:
        return function in self.FUNCTIONS

    def submit(self, job_input: dict, execution_timeout_s: int, function: str = "run_simulation") -> dict:
        """POST /run. Returns {"id": ..., "status": "IN_QUEUE"}."""
        if not self.supports(function):
            raise RunPodError(f"{function} is not available on RunPod (use SIM_BACKEND=modal)")
        payload = {
            "input": {**job_input, "function": function},
            "policy": {"executionTimeout": int(execution_timeout_s * 1000)},
        }
        try:
            r = self._client.post("/run", json=payload)
        except httpx.HTTPError as e:
            raise RunPodError(f"RunPod unreachable: {e}") from e
        if r.status_code != 200:
            raise RunPodError(f"RunPod /run returned HTTP {r.status_code}: {r.text[:300]}")
        data = r.json()
        if "id" not in data:
            raise RunPodError(f"unexpected RunPod response: {str(data)[:300]}")
        return data

    def status(self, job_id: str) -> dict:
        """GET /status/{id}: {"status", "output"?, "error"?, "executionTime"?, "delayTime"?}."""
        try:
            r = self._client.get(f"/status/{job_id}")
        except httpx.HTTPError as e:
            raise RunPodError(f"RunPod unreachable: {e}") from e
        if r.status_code != 200:
            raise RunPodError(f"RunPod /status returned HTTP {r.status_code}: {r.text[:300]}")
        return r.json()