# src/grailrobotics/client.py
import os
import requests
from dataclasses import dataclass
from typing import Any, Dict, Optional, List

DEFAULT_BASE_URL = os.environ.get(
    "GRAILROBOTICS_BASE_URL",
    "https://fly-health.fly.dev",
)

class APIError(RuntimeError):
    pass

def _check(r: requests.Response) -> None:
    if r.ok:
        return
    raise APIError(f"HTTP {r.status_code}: {r.text[:2000]}")

@dataclass
class Client:
    base_url: str = DEFAULT_BASE_URL
    timeout: float = 60.0

    def _url(self, path: str) -> str:
        return self.base_url.rstrip("/") + path

    def health(self) -> Dict[str, Any]:
        r = requests.get(self._url("/health"), timeout=self.timeout)
        _check(r)
        return r.json()

    def start_job(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        r = requests.post(self._url("/jobs"), json=payload, timeout=self.timeout)
        _check(r)
        return r.json()

    def job_status(self, job_id: str) -> Dict[str, Any]:
        r = requests.get(self._url(f"/jobs/{job_id}"), timeout=self.timeout)
        _check(r)
        return r.json()

    def job_logs(self, job_id: str, tail: int = 200) -> Dict[str, Any]:
        r = requests.get(self._url(f"/jobs/{job_id}/logs"), params={"tail": tail}, timeout=self.timeout)
        _check(r)
        return r.json()

    def job_progress(self, job_id: str) -> Dict[str, Any]:
        r = requests.get(self._url(f"/jobs/{job_id}/progress"), timeout=self.timeout)
        _check(r)
        return r.json()

    def cancel_job(self, job_id: str, lambda_api_key: str) -> Dict[str, Any]:
        r = requests.post(
            self._url(f"/jobs/{job_id}/cancel"),
            headers={"X-Lambda-Api-Key": lambda_api_key},
            timeout=self.timeout,
        )
        _check(r)
        return r.json()

    def list_instance_types(self, lambda_api_key: str) -> Dict[str, Any]:
        r = requests.get(
            self._url("/lambda/instance-types"),
            headers={"X-Lambda-Api-Key": lambda_api_key},
            timeout=self.timeout,
        )
        _check(r)
        return r.json()
