"""DataRobot client stub (control plane).

Workload/agentic operations are guarded behind WORKLOAD_API_ENABLED because
those are premium, entitlement-gated features. Set the flag ONLY if your tenant
has the corresponding entitlement.
"""
import os
import requests

BASE = os.getenv("DATAROBOT_ENDPOINT", "https://app.datarobot.com/api/v2")
TOKEN = os.getenv("DATAROBOT_API_TOKEN")
WORKLOADS_ENABLED = os.getenv("WORKLOAD_API_ENABLED", "false").lower() == "true"


def _headers():
    if not TOKEN:
        raise RuntimeError("DATAROBOT_API_TOKEN not set")
    return {"Authorization": f"Bearer {TOKEN}"}


# --- Projects / predictive AI ---
def list_projects():
    return requests.get(f"{BASE}/projects/", headers=_headers(), timeout=30).json()


def get_project(project_id: str):
    return requests.get(f"{BASE}/projects/{project_id}/", headers=_headers(), timeout=30).json()


# --- Deployments (MLOps prediction) ---
def list_deployments():
    return requests.get(f"{BASE}/deployments/", headers=_headers(), timeout=30).json()


def predict(deployment_id: str, data):
    # Batch/prediction endpoints; see registry for exact paths
    return requests.post(
        f"{BASE}/predApi/deployments/{deployment_id}/predictions",
        headers=_headers(), json=data, timeout=60,
    ).json()


# --- Workloads (premium, guarded) ---
def list_workloads():
    if not WORKLOADS_ENABLED:
        raise RuntimeError(
            "Workload API is premium/entitlement-gated; set WORKLOAD_API_ENABLED=true only if entitled."
        )
    return requests.get(f"{BASE}/workloads", headers=_headers(), timeout=30).json()


def create_workload(payload: dict):
    if not WORKLOADS_ENABLED:
        raise RuntimeError("Workload API not enabled for this tenant")
    return requests.post(f"{BASE}/workloads", headers=_headers(), json=payload, timeout=30).json()