from fastapi import FastAPI
import subprocess

app = FastAPI(title="Workspace API")


@app.get("/")
def root():
    return {"message": "workspace api online"}


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/gpu")
def gpu_info():
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,memory.total,driver_version",
                "--format=csv,noheader",
            ],
            text=True,
            timeout=5,
        )
        gpus = [line.strip() for line in out.splitlines() if line.strip()]
        return {"gpus": gpus}
    except Exception as exc:
        return {"gpus": [], "error": str(exc)}