api.py

# api.py
import os
import re
import tempfile
from datetime import datetime
from typing import Any, Dict, Optional

import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field


# ----------------------------
# Configuration
# ----------------------------
ROBOT_HOST = os.getenv("ROBOT_HOST", "127.0.0.1")
ROBOT_PORT_OT2 = int(os.getenv("ROBOT_PORT_OT2", "31950"))
ROBOT_PORT_FLEX = int(os.getenv("ROBOT_PORT_FLEX", "31951"))

OT_API_VERSION = os.getenv("OT_API_VERSION", "3")
HTTP_TIMEOUT = float(os.getenv("HTTP_TIMEOUT", "30"))

# Max bytes of non-JSON text to include in errors
MAX_TEXT_SNIP = int(os.getenv("MAX_TEXT_SNIP", "2000"))


# ----------------------------
# FastAPI app
# ----------------------------
app = FastAPI(title="Opentrons Protocol Analyzer API", version="1.0.0")


# ----------------------------
# Models
# ----------------------------
class AnalyzeRequest(BaseModel):
    protocol_text: str = Field(..., description="Python protocol source code")
    filename: str = Field("protocol.py", description="Filename to send to robot-server")


class AnalyzeResponse(BaseModel):
    ok: bool
    robot_type: str  # "ot2" or "flex"
    robot_server_url: str
    status_code: int
    payload: Dict[str, Any]


# ----------------------------
# Helpers
# ----------------------------
_REQUIREMENTS_ROBOTTYPE_RE = re.compile(
    r"""requirements\s*=\s*{.*?["']robotType["']\s*:\s*["']([^"']+)["'].*?}""",
    re.DOTALL,
)


def detect_robot_type(protocol_text: str) -> str:
    """
    Detect robot type from protocol 'requirements' metadata.

    Returns:
      "flex" if requirements.robotType == "Flex"
      otherwise "ot2" (default)
    """
    m = _REQUIREMENTS_ROBOTTYPE_RE.search(protocol_text)
    if not m:
        return "ot2"

    val = (m.group(1) or "").strip()
    if val == "Flex":
        return "flex"
    if val == "OT-2":
        return "ot2"

    # Unknown value -> be conservative
    return "ot2"


def robot_server_url(robot_type: str) -> str:
    port = ROBOT_PORT_FLEX if robot_type == "flex" else ROBOT_PORT_OT2
    return f"http://{ROBOT_HOST}:{port}"


def post_protocol(protocol_text: str, filename: str, base_url: str) -> requests.Response:
    endpoint = f"{base_url}/protocols"
    headers = {"Opentrons-Version": OT_API_VERSION}

    fd, path = tempfile.mkstemp(prefix="protocol_", suffix=".py")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(protocol_text)

        with open(path, "rb") as f:
            files = {"files": (filename, f, "text/x-python")}
            return requests.post(endpoint, headers=headers, files=files, timeout=HTTP_TIMEOUT)
    finally:
        try:
            os.remove(path)
        except OSError:
            pass


def safe_json(resp: requests.Response) -> Dict[str, Any]:
    try:
        data = resp.json()
        if isinstance(data, dict):
            return data
        return {"_non_dict_json": data}
    except ValueError:
        return {"_non_json": resp.text[:MAX_TEXT_SNIP]}


# ----------------------------
# Routes
# ----------------------------
APP_NAME = os.getenv("APP_NAME", "Opentrons Protocol Analyzer API")
APP_VERSION = os.getenv("APP_VERSION", "dev")
BUILD_SHA = os.getenv("BUILD_SHA", "unknown")

@app.get("/")
def root():
    return {
        "name": APP_NAME,
        "version": APP_VERSION,
        "build": BUILD_SHA,
        "status": "ok",
        "time": datetime.utcnow().isoformat() + "Z",
    }

@app.get("/health")
def health() -> Dict[str, Any]:
    return {"ok": True}


@app.get("/robot/health")
def robot_health() -> dict:
    out: dict = {"ok": True, "robots": {}}

    headers = {"Opentrons-Version": OT_API_VERSION}

    for rt in ("ot2", "flex"):
        base = robot_server_url(rt)
        try:
            r = requests.get(f"{base}/health", headers=headers, timeout=5)
            out["robots"][rt] = {"ok": r.ok, "status_code": r.status_code, "text": r.text[:2000]}
            if not r.ok:
                out["ok"] = False
        except requests.RequestException as e:
            out["robots"][rt] = {"ok": False, "error": str(e)}
            out["ok"] = False
    return out



@app.post("/analyze", response_model=AnalyzeResponse)
def analyze(req: AnalyzeRequest) -> AnalyzeResponse:
    protocol_text = (req.protocol_text or "").strip()
    if not protocol_text:
        raise HTTPException(status_code=400, detail="protocol_text is empty")

    rt = detect_robot_type(protocol_text)
    base = robot_server_url(rt)

    try:
        resp = post_protocol(protocol_text, req.filename, base)
    except requests.RequestException as e:
        raise HTTPException(status_code=503, detail=f"robot-server request failed: {e}")

    payload = safe_json(resp)

    return AnalyzeResponse(
        ok=resp.status_code < 400,
        robot_type=rt,
        robot_server_url=base,
        status_code=resp.status_code,
        payload=payload,
    )