cx-ui/backend/app/routers/runpod_router.py

"""RunPod machines: current state, per-host history, and problem ranking."""
from __future__ import annotations

import datetime as dt

from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel
from sqlalchemy import func, select
from sqlalchemy.orm import Session

from ..auth import current_user
from ..config import get_settings
from ..db import get_db
from ..models import RunpodColour, RunpodEvent, RunpodEventType, RunpodHost, User
from ..runpod.client import RunPodClient, RunPodError
from ..runpod.email_parse import parse as parse_email
from ..runpod.delivery import raise_rma_issue, raise_unlisting_ticket, unlisting_ticket
from ..runpod.service import record_event, sync_machines
from ..delivery import DeliveryError

router = APIRouter(prefix="/api/runpod", tags=["runpod"])
settings = get_settings()


def _client() -> RunPodClient:
    return RunPodClient(
        api_key=settings.runpod_api_key, email=settings.runpod_email,
        password=settings.runpod_password, team_id=settings.runpod_team_id,
        totp_secret=settings.runpod_totp_secret,
    )


class HostPatch(BaseModel):
    colour: str | None = None
    zendesk_ticket: str | None = None
    jira_key: str | None = None
    jira_status: str | None = None
    next_steps: str | None = None
    last_error: str | None = None


class NoteBody(BaseModel):
    detail: str
    event_type: str = "note"
    zendesk_ticket: str = ""
    jira_key: str = ""


class EmailBody(BaseModel):
    raw: str
    subject: str = ""


class TicketBody(BaseModel):
    error_text: str = ""
    requester: str = ""
    subject: str = ""
    body: str = ""


class RmaBody(BaseModel):
    summary: str = ""
    description: str = ""


@router.get("/status")
def runpod_status(db: Session = Depends(get_db), user: User = Depends(current_user)):
    hosts = db.scalars(select(RunpodHost)).all()
    unlisted = [h for h in hosts if not h.listed]
    return {
        "configured": settings.runpod_ready,
        "mode": _client().mode,
        "write_enabled": settings.feature_runpod_write,
        "totals": {
            "machines": len(hosts),
            "listed": len(hosts) - len(unlisted),
            "unlisted": len(unlisted),
            "gpus_rented": sum(h.gpu_reserved for h in hosts),
            "gpus_total": sum(h.gpu_total for h in hosts),
        },
        "colours": [c.value for c in RunpodColour],
        "event_types": [e.value for e in RunpodEventType],
    }


@router.get("/hosts")
def hosts(unlisted_only: bool = False, db: Session = Depends(get_db),
          user: User = Depends(current_user)):
    query = select(RunpodHost)
    if unlisted_only:
        query = query.where(RunpodHost.listed.is_(False))
    rows = db.scalars(query.order_by(RunpodHost.listed, RunpodHost.name)).all()
    return {"hosts": [h.to_json() for h in rows]}


@router.get("/hosts/{machine_id}")
def host_detail(machine_id: str, db: Session = Depends(get_db), user: User = Depends(current_user)):
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    return row.to_json(with_events=True)


@router.get("/problem-hosts")
def problem_hosts(limit: int = 40, db: Session = Depends(get_db),
                  user: User = Depends(current_user)):
    """Machines ranked by how often they have been unlisted.

    The repeat offenders are the ones worth an RMA conversation rather than
    another burn-in, which is what the old script's historical table showed.
    """
    unlist_counts = (
        select(RunpodEvent.host_id, func.count(RunpodEvent.id).label("n"),
               func.max(RunpodEvent.occurred_at).label("last"))
        .where(RunpodEvent.event_type == RunpodEventType.UNLISTED)
        .group_by(RunpodEvent.host_id).subquery()
    )
    rows = db.execute(
        select(RunpodHost, unlist_counts.c.n, unlist_counts.c.last)
        .join(unlist_counts, unlist_counts.c.host_id == RunpodHost.id, isouter=True)
        .order_by(func.coalesce(unlist_counts.c.n, 0).desc(), RunpodHost.historic_count.desc())
        .limit(max(1, min(limit, 200)))
    ).all()

    out = []
    for host, count, last in rows:
        total = int(count or 0) or host.historic_count
        if not total:
            continue
        payload = host.to_json()
        payload["unlist_events"] = int(count or 0)
        payload["effective_count"] = total
        payload["last_unlisted"] = last.isoformat() if last else (
            host.unlisted_at.isoformat() if host.unlisted_at else None)
        out.append(payload)
    return {"hosts": out}


@router.post("/sync")
def sync(db: Session = Depends(get_db), user: User = Depends(current_user)):
    if not settings.runpod_ready:
        raise HTTPException(status.HTTP_400_BAD_REQUEST,
                            "RunPod is not configured. Set CX_RUNPOD_API_KEY.")
    try:
        result = sync_machines(db, _client(), actor=user.email)
    except RunPodError as exc:
        raise HTTPException(status.HTTP_502_BAD_GATEWAY, str(exc)) from exc
    return result


@router.post("/hosts/{machine_id}")
def patch_host(machine_id: str, body: HostPatch, db: Session = Depends(get_db),
               user: User = Depends(current_user)):
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    changes = []
    if body.colour is not None:
        try:
            row.colour = RunpodColour(body.colour)
        except ValueError as exc:
            raise HTTPException(status.HTTP_400_BAD_REQUEST, f"Unknown colour '{body.colour}'") from exc
        changes.append(f"colour={body.colour}")
    for field in ("zendesk_ticket", "jira_key", "jira_status", "next_steps", "last_error"):
        value = getattr(body, field)
        if value is not None:
            setattr(row, field, value)
            changes.append(field)
    if changes:
        record_event(db, row, RunpodEventType.NOTE, actor=user.email,
                     detail="Updated " + ", ".join(changes))
    db.commit()
    return row.to_json(with_events=True)


@router.post("/hosts/{machine_id}/events")
def add_event(machine_id: str, body: NoteBody, db: Session = Depends(get_db),
              user: User = Depends(current_user)):
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    try:
        event_type = RunpodEventType(body.event_type)
    except ValueError as exc:
        raise HTTPException(status.HTTP_400_BAD_REQUEST, f"Unknown event type") from exc
    record_event(db, row, event_type, actor=user.email, detail=body.detail,
                 zendesk_ticket=body.zendesk_ticket, jira_key=body.jira_key)
    db.commit()
    return row.to_json(with_events=True)


@router.get("/hosts/{machine_id}/ticket-preview")
def ticket_preview(machine_id: str, db: Session = Depends(get_db),
                   user: User = Depends(current_user)):
    """The Zendesk ticket that would be raised, without raising it."""
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    return {
        **unlisting_ticket(row),
        "zendesk_ready": settings.zendesk_ready,
        "send_enabled": settings.feature_send_enabled,
        "runpod_jira_ready": settings.runpod_jira_ready,
        "existing_ticket": row.zendesk_ticket,
        "existing_jira": row.jira_key,
    }


@router.post("/hosts/{machine_id}/zendesk")
async def zendesk_ticket(machine_id: str, body: TicketBody, db: Session = Depends(get_db),
                         user: User = Depends(current_user)):
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    try:
        return await raise_unlisting_ticket(db, row, user, error_text=body.error_text,
                                            requester=body.requester, subject=body.subject,
                                            body=body.body)
    except DeliveryError as exc:
        raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc


@router.post("/hosts/{machine_id}/rma")
async def rma_issue(machine_id: str, body: RmaBody, db: Session = Depends(get_db),
                    user: User = Depends(current_user)):
    row = db.scalars(select(RunpodHost).where(RunpodHost.machine_id == machine_id)).first()
    if row is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such machine")
    try:
        return await raise_rma_issue(db, row, user, summary=body.summary,
                                     description=body.description)
    except DeliveryError as exc:
        raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc


@router.post("/ingest-email")
def ingest_email(body: EmailBody, db: Session = Depends(get_db), user: User = Depends(current_user)):
    """Take a RunPod unlisting email and attach its error hint to the machine.

    The email is the only place the actual failure reason appears, so it is
    worth capturing even when the unlisting itself came through the API.
    """
    parsed = parse_email(body.raw, body.subject)
    if not parsed["machine_id"] and not parsed["host"]:
        raise HTTPException(status.HTTP_400_BAD_REQUEST,
                            "Could not find a machine in that email.")
    query = select(RunpodHost)
    row = db.scalars(query.where(RunpodHost.machine_id == parsed["machine_id"])).first() \
        if parsed["machine_id"] else None
    if row is None and parsed["host"]:
        row = db.scalars(query.where(RunpodHost.name == parsed["host"])).first()
    if row is None:
        return {"parsed": parsed, "matched": False,
                "detail": "Parsed the email, but no machine in the database matches."}

    row.last_error = parsed["error_text"] or row.last_error
    record_event(db, row, RunpodEventType.UNLISTED, actor="runpod",
                 detail=f"{parsed['signature']} — {parsed['suggested_action']}",
                 error_hint=parsed["error_text"], gpu_reserved=parsed["gpus_rented"],
                 payload={"source": "email", "critical": parsed["is_critical"]})
    db.commit()
    return {"parsed": parsed, "matched": True, "host": row.to_json(with_events=True)}