periscope/backend/routers/pipeline.py

"""Pipeline start, SSE events, and status endpoints.

Pipelines run in a Cloud Run Job worker (or, in local dev, a child
subprocess). The API only enqueues, transitions status with
``if-generation-match`` for idempotency, and tails the GCS-backed event
log for SSE.
"""

from __future__ import annotations

import asyncio
import json
import logging

from fastapi import APIRouter, HTTPException, Request
from sse_starlette.sse import EventSourceResponse

from pydantic import BaseModel

from backend.routers.deps import get_storage, resolve_or_404
from backend.services import event_bridge, job_runner
from backend.services import projects as proj_svc

logger = logging.getLogger(__name__)

VALID_REGEN_STAGES = {"derating"}


class RegenRequest(BaseModel):
    stages: list[str]


router = APIRouter(tags=["pipeline"])


# Statuses from which a fresh ``/start`` is allowed to transition into queued.
_START_OK_FROM = frozenset({
    proj_svc.STATUS_DRAFT,
    proj_svc.STATUS_COMPLETE,
    proj_svc.STATUS_ERROR,
    proj_svc.STATUS_CANCELLED,
})


def _project_active(meta: proj_svc.ProjectMeta) -> bool:
    """A project is "active" if a worker is or could be running for it.

    Used as the running-guard. We trust the meta status as the primary
    signal, and only fall back to the Cloud Run execution state when the
    status is one we expect a worker to be touching. This deliberately
    does NOT call get_execution_state on every request — it's an admin
    API call. The stale-running sweeper is responsible for clearing
    zombie ``running`` projects.
    """
    return meta.status in (proj_svc.STATUS_QUEUED, proj_svc.STATUS_RUNNING)


@router.post("/pipeline/{project_id}/start", status_code=202)
async def start(project_id: str, request: Request):
    from backend.routers.deps import get_user_id
    from backend.services.billing_hook import get_billing

    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if not meta.has_bom or not meta.has_netlist:
        raise HTTPException(400, "Upload BOM and netlist before starting pipeline")

    # Ensure the caller has at least their trial credits allocated.  The
    # pipeline itself enforces pause-on-empty — this just makes sure a
    # brand-new user isn't blocked before their grant is issued.
    get_billing().ensure_trial_grant(storage, get_user_id(request))

    # Idempotent enqueue: only one ``draft|complete|error|cancelled`` ->
    # ``queued`` transition can win. Concurrent /start clicks => 409.
    from backend._version import PINSCOPE_VERSION
    try:
        proj_svc.transition_status(
            storage, owner_user_id, project_id,
            from_status=_START_OK_FROM,
            to_status=proj_svc.STATUS_QUEUED,
            cancel_requested=False,
            execution_name=None,
            pinscope_version=PINSCOPE_VERSION,
        )
    except proj_svc.StatusConflict:
        raise HTTPException(409, "Pipeline already running or queued")

    try:
        execution_name = job_runner.enqueue_pipeline(
            project_id, owner_user_id, resume=False, free=False,
        )
    except Exception:
        logger.exception("enqueue_pipeline failed for %s", project_id)
        # Roll the meta back so the user can retry.
        proj_svc.update_project(
            storage, owner_user_id, project_id,
            status=proj_svc.STATUS_ERROR,
            pipeline_state={"error": "Failed to enqueue worker"},
        )
        raise HTTPException(503, "Failed to enqueue pipeline worker; please retry")

    proj_svc.update_project(
        storage, owner_user_id, project_id, execution_name=execution_name,
    )
    return {"status": "started", "project_id": project_id}


@router.post("/pipeline/{project_id}/cancel")
async def cancel(project_id: str, request: Request):
    """Soft-cancel: set ``cancel_requested`` so the worker exits cleanly.

    The worker re-reads this flag inside ``_charge_for_logs`` after every
    Claude API call (throttled). Cancellation latency is bounded by the
    in-flight call's duration, typ 1–60s.
    """
    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if not _project_active(meta):
        raise HTTPException(409, f"Pipeline is not running (status={meta.status})")
    proj_svc.request_cancel(storage, owner_user_id, project_id)
    return {"status": "cancel_requested", "project_id": project_id}


@router.post("/pipeline/{project_id}/estimate")
async def estimate(project_id: str, request: Request):
    """Pre-flight cost estimate — read-only, no side effects."""
    from backend.services.cost_estimator import estimate_pipeline_cost

    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if not meta.has_bom:
        raise HTTPException(400, "Upload a BOM before requesting an estimate")
    try:
        est = estimate_pipeline_cost(storage, owner_user_id, project_id)
    except FileNotFoundError as exc:
        raise HTTPException(400, str(exc)) from exc
    return est.model_dump()


@router.post("/pipeline/{project_id}/resume", status_code=202)
async def resume(project_id: str, request: Request):
    """Resume a pipeline that was paused for insufficient credits."""
    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if meta.status != proj_svc.STATUS_PAUSED:
        raise HTTPException(
            400,
            f"Project is not paused (status={meta.status}); nothing to resume.",
        )
    if not meta.has_bom or not meta.has_netlist:
        raise HTTPException(400, "Project is missing BOM or netlist")

    try:
        proj_svc.transition_status(
            storage, owner_user_id, project_id,
            from_status=proj_svc.STATUS_PAUSED,
            to_status=proj_svc.STATUS_QUEUED,
            cancel_requested=False,
        )
    except proj_svc.StatusConflict:
        raise HTTPException(409, "Project state changed; refresh and retry")

    try:
        execution_name = job_runner.enqueue_pipeline(
            project_id, owner_user_id, resume=True, free=False,
        )
    except Exception:
        logger.exception("enqueue_pipeline (resume) failed for %s", project_id)
        proj_svc.update_project(
            storage, owner_user_id, project_id,
            status=proj_svc.STATUS_ERROR,
            pipeline_state={"error": "Failed to enqueue worker"},
        )
        raise HTTPException(503, "Failed to enqueue pipeline worker; please retry")

    proj_svc.update_project(
        storage, owner_user_id, project_id, execution_name=execution_name,
    )
    return {"status": "resumed", "project_id": project_id}


@router.post("/pipeline/{project_id}/restart", status_code=202)
async def restart(project_id: str, request: Request):
    """Admin-only: wipe per-project extractions and run the pipeline free."""
    from backend.routers.admin import _require_admin

    await _require_admin(request)
    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if not meta.has_bom or not meta.has_netlist:
        raise HTTPException(400, "Upload BOM and netlist before starting pipeline")

    # If something is currently running/queued, request cancel and wait
    # briefly for the worker to honour it (or exit on its own). Hard-kill
    # the execution as a last resort.
    if _project_active(meta):
        proj_svc.request_cancel(storage, owner_user_id, project_id)
        await _await_terminal(storage, owner_user_id, project_id, timeout_s=10.0)
        # If still active, hard-kill via Cloud Run cancel.
        meta = proj_svc.get_project(storage, owner_user_id, project_id) or meta
        if _project_active(meta) and meta.execution_name:
            job_runner.cancel_execution(meta.execution_name)
            await _await_terminal(storage, owner_user_id, project_id, timeout_s=5.0)

    proj_svc.clear_project_extractions(storage, owner_user_id, project_id)

    # After clear_project_extractions the project is left in whatever
    # status it was; the transition below enforces queued.
    try:
        proj_svc.transition_status(
            storage, owner_user_id, project_id,
            from_status=_START_OK_FROM | {proj_svc.STATUS_PAUSED},
            to_status=proj_svc.STATUS_QUEUED,
            cancel_requested=False,
            execution_name=None,
        )
    except proj_svc.StatusConflict:
        raise HTTPException(409, "Pipeline is busy; cancel first then retry")

    try:
        execution_name = job_runner.enqueue_pipeline(
            project_id, owner_user_id, resume=False, free=True,
        )
    except Exception:
        logger.exception("enqueue_pipeline (restart) failed for %s", project_id)
        proj_svc.update_project(
            storage, owner_user_id, project_id,
            status=proj_svc.STATUS_ERROR,
            pipeline_state={"error": "Failed to enqueue worker"},
        )
        raise HTTPException(503, "Failed to enqueue pipeline worker; please retry")

    proj_svc.update_project(
        storage, owner_user_id, project_id, execution_name=execution_name,
    )
    return {"status": "restarted", "project_id": project_id}


@router.post("/pipeline/{project_id}/regen", status_code=202)
async def regen(project_id: str, req: RegenRequest, request: Request):
    """Rebuild graph and regenerate only the requested stages."""
    storage = get_storage(request)
    owner_user_id, meta = await resolve_or_404(request, project_id)
    if not meta.has_bom or not meta.has_netlist:
        raise HTTPException(400, "Upload BOM and netlist before running regen")
    invalid = set(req.stages) - VALID_REGEN_STAGES
    if invalid:
        raise HTTPException(400, f"Invalid regen stages: {sorted(invalid)}. Valid: {sorted(VALID_REGEN_STAGES)}")
    if not req.stages:
        raise HTTPException(400, "At least one stage is required")

    if _project_active(meta):
        proj_svc.request_cancel(storage, owner_user_id, project_id)
        await _await_terminal(storage, owner_user_id, project_id, timeout_s=10.0)
        meta = proj_svc.get_project(storage, owner_user_id, project_id) or meta
        if _project_active(meta) and meta.execution_name:
            job_runner.cancel_execution(meta.execution_name)
            await _await_terminal(storage, owner_user_id, project_id, timeout_s=5.0)

    try:
        proj_svc.transition_status(
            storage, owner_user_id, project_id,
            from_status=_START_OK_FROM | {proj_svc.STATUS_PAUSED},
            to_status=proj_svc.STATUS_QUEUED,
            cancel_requested=False,
            execution_name=None,
        )
    except proj_svc.StatusConflict:
        raise HTTPException(409, "Pipeline is busy; cancel first then retry")

    try:
        execution_name = job_runner.enqueue_pipeline_regen(
            project_id, owner_user_id, stages=req.stages,
        )
    except Exception:
        logger.exception("enqueue_pipeline_regen failed for %s", project_id)
        proj_svc.update_project(
            storage, owner_user_id, project_id,
            status=proj_svc.STATUS_ERROR,
            pipeline_state={"error": "Failed to enqueue worker"},
        )
        raise HTTPException(503, "Failed to enqueue pipeline worker; please retry")

    proj_svc.update_project(
        storage, owner_user_id, project_id, execution_name=execution_name,
    )
    return {"status": "regen_started", "project_id": project_id, "stages": req.stages}


# ---------------------------------------------------------------------------
# SSE events
# ---------------------------------------------------------------------------


_EXEC_TERMINAL = frozenset({"succeeded", "failed", "cancelled"})


@router.get("/pipeline/{project_id}/events")
async def events(project_id: str, request: Request):
    """SSE stream of pipeline progress events.

    Tails the GCS-backed event log written by the worker. Stops on
    terminal events as today, but also has two hard-crash escape
    hatches: the project's status reaching a terminal value, and the
    Cloud Run execution reaching a terminal state. Either of those
    triggers a synthetic ``pipeline_error`` so the SSE doesn't hang
    forever when the worker dies without writing its terminal event.
    """
    owner_user_id, meta = await resolve_or_404(request, project_id)
    storage = get_storage(request)

    async def event_generator():
        execution_name = meta.execution_name
        # Drive the GCS tail and the escape-hatch poll concurrently. The
        # tail yields events; the escape hatch flips a flag.
        crash_detected: dict[str, str | None] = {"reason": None}

        async def watch_status() -> None:
            poll_interval = 2.0
            while True:
                await asyncio.sleep(poll_interval)
                try:
                    cur = proj_svc.get_project(storage, owner_user_id, project_id)
                except Exception:
                    continue
                if cur is None:
                    continue
                if cur.status in proj_svc.TERMINAL_STATUSES:
                    crash_detected["reason"] = (
                        f"project status={cur.status} (terminal)"
                    )
                    return
                # Cloud Run hard-crash detection
                if execution_name:
                    try:
                        state = job_runner.get_execution_state(execution_name)
                    except Exception:
                        state = "unknown"
                    if state in _EXEC_TERMINAL:
                        crash_detected["reason"] = (
                            f"execution state={state}"
                        )
                        return

        watcher = asyncio.create_task(watch_status())
        try:
            async for msg in event_bridge.tail_events(
                storage, owner_user_id, project_id,
            ):
                if crash_detected["reason"] is not None:
                    break
                yield {
                    "event": msg["event"],
                    "data": json.dumps(msg.get("data", {})),
                }
                if msg["event"] in event_bridge.TERMINAL_EVENTS:
                    return

            # tail_events exited without a terminal event — escape hatch
            if crash_detected["reason"] is not None:
                # Re-read the current meta so the synthetic event has
                # the most up-to-date error information.
                cur = proj_svc.get_project(storage, owner_user_id, project_id)
                err = (
                    (cur.pipeline_state or {}).get("error")
                    if cur and cur.pipeline_state
                    else crash_detected["reason"]
                )
                yield {
                    "event": "pipeline_error",
                    "data": json.dumps({
                        "error": err or "worker terminated without writing a terminal event",
                        "synthetic": True,
                    }),
                }
        finally:
            watcher.cancel()
            try:
                await watcher
            except (asyncio.CancelledError, Exception):
                pass

    return EventSourceResponse(event_generator())


@router.get("/pipeline/{project_id}/status")
async def status(project_id: str, request: Request):
    """Polling fallback — returns current project state."""
    _, meta = await resolve_or_404(request, project_id)
    return {
        "status": meta.status,
        "summary": meta.summary,
        "pipeline_state": meta.pipeline_state,
        "running": meta.status in (proj_svc.STATUS_RUNNING, proj_svc.STATUS_QUEUED),
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


async def _await_terminal(
    storage, user_id: str, project_id: str, *, timeout_s: float,
) -> None:
    """Poll project status until it reaches a terminal state or the timeout
    elapses. Used by /restart and /regen between cancel and re-enqueue.
    """
    poll = 0.5
    elapsed = 0.0
    while elapsed < timeout_s:
        try:
            meta = proj_svc.get_project(storage, user_id, project_id)
        except Exception:
            meta = None
        if meta is None:
            return
        if meta.status in proj_svc.TERMINAL_STATUSES:
            return
        await asyncio.sleep(poll)
        elapsed += poll