app/backend/app/api/routes/flows.py

"""The flow API. Everything the editor can do is available here first."""

import asyncio
import time
from typing import Any

from fastapi import (
    APIRouter,
    Depends,
    HTTPException,
    WebSocket,
    WebSocketDisconnect,
)
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel
from sqlmodel import Session

from app.api.deps import (
    CurrentUser,
    FlowControllerDep,
    get_current_user,
    user_from_token,
)
from app.core.db import engine
from app.flow.controller import FlowController
from app.flow.dashboards import DashboardStore
from app.flow.events import event_bus
from app.flow.messages import qualify
from app.flow.pipeline import ValidationIssue
from app.flow.schemas import (
    NAME_PATTERN,
    BrainGraph,
    FlowDef,
    FlowsPublic,
    FlowStatePublic,
    FlowSummary,
    HistoryPoint,
    LibraryNode,
    MessageHistory,
    MessageValue,
    NodeSource,
    NodeStatusPublic,
    NodeTypeInfo,
)
from app.flow.state import as_number
from app.flow.store import (
    FlowExists,
    FlowNotFound,
    LibExists,
    LibNotFound,
    StaleVersion,
)
from app.models import Message

router = APIRouter(
    prefix="/flows", tags=["flows"], dependencies=[Depends(get_current_user)]
)
# The websocket authenticates from its query string, so it stays off the
# router that enforces the HTTP bearer scheme.
ws_router = APIRouter(prefix="/flows", tags=["flows"])


class Endpoint(BaseModel):
    """Something wired into this flow that is not a node in it.

    A dashboard control setting one of its messages, a tile showing one, or a
    node in another flow on the far side of a dotted name. The canvas draws
    these so a value never appears to come from nowhere — or worse, appears to
    come from whichever node happens to be drawn as a producer.
    """

    #: dashboard or flow.
    kind: str
    #: Stable within its kind, and used as the canvas node id.
    id: str
    label: str
    #: What sort of widget, or the node type in the other flow.
    detail: str = ""
    #: Messages of this flow it publishes, and ones it reads.
    provides: list[str] = []
    requires: list[str] = []


class FlowDetail(BaseModel):
    """A flow plus how it is currently doing.

    ``definition`` is the working copy — the unpublished draft when there is
    one — because that is what the editor shows. ``nodes`` reports the
    published flow, which is what is actually running.
    """

    definition: FlowDef
    nodes: list[NodeStatusPublic] = []
    issues: list[ValidationIssue] = []
    has_draft: bool = False
    enabled: bool = True
    paused: bool = False
    #: Dashboards and other flows wired into this one.
    endpoints: list[Endpoint] = []


class ValidationResult(BaseModel):
    issues: list[ValidationIssue] = []


class RenameRequest(BaseModel):
    new_name: str


class PublishRequest(BaseModel):
    version: int


class ShareRequest(BaseModel):
    lib_name: str


class RunRequest(BaseModel):
    inputs: dict[str, Any] = {}


class TriggerRequest(BaseModel):
    values: dict[str, Any] = {}


def _endpoints(controller: FlowController, flow: str) -> list[Endpoint]:
    """Everything wired into ``flow`` from outside it."""
    found: list[Endpoint] = []

    dashboards: DashboardStore | None = getattr(controller, "dashboards", None)
    if dashboards is not None:
        for binding in dashboards.bindings_for(flow):
            found.append(
                Endpoint(
                    kind="dashboard",
                    id=f"dashboard:{binding['dashboard']}:{binding['widget']}",
                    label=binding["title"],
                    detail=binding["type"],
                    provides=[binding["provides"]] if binding["provides"] else [],
                    requires=binding["requires"],
                )
            )

    for other, node_id, provides, requires in controller.cross_flow_nodes(flow):
        found.append(
            Endpoint(
                kind="flow",
                id=f"flow:{other}.{node_id}",
                label=f"{other}.{node_id}",
                detail="flow",
                provides=provides,
                requires=requires,
            )
        )
    return found


def _detail(controller: FlowController, definition: FlowDef) -> FlowDetail:
    name = definition.name
    running = {
        "enabled": controller.is_enabled(name),
        "paused": controller.is_paused(name),
    }
    endpoints = _endpoints(controller, name)
    if controller.store.has_draft(name):
        # Report the draft the editor is showing, not the version running
        # underneath it — otherwise a node the author just broke looks fine.
        preview = controller.preview(name)
        return FlowDetail(
            definition=definition,
            nodes=preview.nodes,
            issues=preview.issues,
            has_draft=True,
            endpoints=endpoints,
            **running,
        )
    return FlowDetail(
        definition=definition,
        nodes=controller.node_statuses(name),
        issues=controller.flow_issues(name),
        endpoints=endpoints,
        **running,
    )


def _read_flow(controller: FlowController, name: str) -> FlowDef:
    """The working copy: the draft when there is one, else what is published."""
    try:
        return controller.store.read_flow(name, draft=True)
    except FlowNotFound:
        raise HTTPException(status_code=404, detail=f"No flow named '{name}'")


def _audit(action: str, flow: str, user: CurrentUser) -> None:
    """Record who changed what. The collector writes it down; the bus carries it."""
    event_bus.publish(
        {
            "type": "audit",
            "action": action,
            "flow": flow,
            "user": user.email,
            "ts": time.time(),
        }
    )


def _source_ref(definition: FlowDef, node_id: str) -> str | None:
    """The library source this node runs, if it is a shared one."""
    node = next((n for n in definition.nodes if n.id == node_id), None)
    return node.source_ref if node else None


def _require_enabled(controller: FlowController, name: str) -> None:
    if not controller.is_enabled(name):
        raise HTTPException(
            status_code=409,
            detail=f"Flow '{name}' is stopped — start it before running it",
        )


def _flow_state(controller: FlowController, name: str) -> FlowStatePublic:
    return FlowStatePublic(
        values={
            key: MessageValue(**value) for key, value in controller.values(name).items()
        },
        nodes=controller.node_statuses(name),
    )


# -----------------------------------------------------------------------------
# Flows
# -----------------------------------------------------------------------------


@router.get("/", response_model=FlowsPublic)
def read_flows(controller: FlowControllerDep) -> Any:
    """List every flow."""
    summaries = []
    for name in controller.store.list_flows():
        try:
            definition = controller.store.read_flow(name, draft=True)
        except FlowNotFound:
            continue
        statuses = controller.node_statuses(name)
        summaries.append(
            FlowSummary(
                name=definition.name,
                title=definition.title,
                node_count=len(definition.nodes),
                error_count=sum(1 for s in statuses if s.status == "error"),
                has_draft=controller.store.has_draft(name),
                enabled=controller.is_enabled(name),
                paused=controller.is_paused(name),
                quarantined=controller.is_quarantined(name),
            )
        )
    return FlowsPublic(data=summaries, count=len(summaries))


@router.get("/node-types", response_model=list[NodeTypeInfo])
def read_node_types() -> Any:
    """The node types that can be placed on a canvas."""
    from app.flow.controller import node_type_info

    return node_type_info()


# Above "/{name}" for the same reason "library" is: a flow called "graph" would
# otherwise be unreachable.
@router.get("/graph", response_model=BrainGraph)
def read_graph(controller: FlowControllerDep) -> Any:
    """Every flow as one graph, with nodes talking to the same thing merged."""
    return controller.brain_graph()


# -----------------------------------------------------------------------------
# Shared nodes
#
# Declared above the "/{name}" routes: "library" would otherwise be read as a
# flow name.
# -----------------------------------------------------------------------------


@router.get("/library", response_model=list[LibraryNode])
def read_library(controller: FlowControllerDep) -> Any:
    """The node sources shared across flows, and which nodes use each."""
    return [
        LibraryNode(name=name, used_by=controller.store.usages(name))
        for name in controller.store.list_lib()
    ]


@router.delete("/library/{lib_name}", response_model=Message)
async def delete_shared_node(lib_name: str, controller: FlowControllerDep) -> Any:
    """Remove a shared source, as long as no flow still runs it."""
    used_by = controller.store.usages(lib_name)
    if used_by:
        raise HTTPException(
            status_code=409,
            detail=f"'{lib_name}' is still used by {', '.join(used_by)}",
        )
    try:
        await run_in_threadpool(controller.store.delete_lib_source, lib_name)
    except LibNotFound:
        raise HTTPException(status_code=404, detail=f"No shared node '{lib_name}'")
    return Message(message=f"Deleted shared node '{lib_name}'")


@router.get("/{name}", response_model=FlowDetail)
def read_flow(name: str, controller: FlowControllerDep) -> Any:
    """Read one flow, with the state of its nodes."""
    return _detail(controller, _read_flow(controller, name))


@router.put("/{name}", response_model=FlowDetail)
async def save_flow(
    name: str,
    definition: FlowDef,
    controller: FlowControllerDep,
) -> Any:
    """Save unpublished changes to a flow.

    This writes a draft: the running pipeline keeps the published version until
    someone publishes. ``version`` is the one the editor last saw — a mismatch
    means another client saved in between and answers 409 rather than throwing
    their work away.
    """
    if definition.name != name:
        raise HTTPException(
            status_code=400, detail="The flow name in the body must match the URL"
        )

    duplicates = {n.id for n in definition.nodes}
    if len(duplicates) != len(definition.nodes):
        raise HTTPException(status_code=400, detail="Node names must be unique")

    try:
        stored = await run_in_threadpool(
            controller.store.write_draft, definition, definition.version
        )
    except StaleVersion as exc:
        raise HTTPException(
            status_code=409,
            detail={"message": str(exc), "current_version": exc.current},
        )
    return _detail(controller, stored)


@router.post("/{name}/publish", response_model=FlowDetail)
async def publish_flow(
    name: str,
    body: PublishRequest,
    controller: FlowControllerDep,
    user: CurrentUser,
) -> Any:
    """Deploy the unpublished changes: the engine picks them up from here."""
    _read_flow(controller, name)
    if not controller.store.has_draft(name):
        raise HTTPException(
            status_code=400, detail=f"Flow '{name}' has no unpublished changes"
        )
    try:
        published = await run_in_threadpool(
            controller.store.publish_flow, name, body.version
        )
    except StaleVersion as exc:
        raise HTTPException(
            status_code=409,
            detail={"message": str(exc), "current_version": exc.current},
        )
    _audit("published", name, user)
    await controller.reload()
    return _detail(controller, published)


@router.post("/{name}/discard-draft", response_model=FlowDetail)
async def discard_draft(name: str, controller: FlowControllerDep) -> Any:
    """Throw the unpublished changes away and go back to what is running."""
    _read_flow(controller, name)
    if not controller.store.has_draft(name):
        raise HTTPException(
            status_code=400, detail=f"Flow '{name}' has no unpublished changes"
        )
    if not controller.store.is_published(name):
        raise HTTPException(
            status_code=400,
            detail=(
                f"Flow '{name}' has never been published — delete it instead of "
                "discarding it"
            ),
        )
    published = await run_in_threadpool(controller.store.discard_draft, name)
    return _detail(controller, published)


@router.delete("/{name}", response_model=Message)
async def delete_flow(
    name: str, controller: FlowControllerDep, user: CurrentUser
) -> Any:
    """Delete a flow and everything in it."""
    try:
        await run_in_threadpool(controller.store.delete_flow, name)
    except FlowNotFound:
        raise HTTPException(status_code=404, detail=f"No flow named '{name}'")
    _audit("deleted", name, user)
    # Its files are gone; its values and queued work would otherwise linger.
    await run_in_threadpool(controller.forget_flow, name)
    await controller.reload()
    return Message(message=f"Deleted flow '{name}'")


@router.post("/{name}/rename", response_model=FlowDetail)
async def rename_flow(
    name: str,
    body: RenameRequest,
    controller: FlowControllerDep,
) -> Any:
    """Rename a flow, along with every reference to its messages."""
    if not NAME_PATTERN.match(body.new_name):
        raise HTTPException(
            status_code=400,
            detail=(
                "Use lowercase letters, digits and underscores, starting with a letter"
            ),
        )
    try:
        renamed = await run_in_threadpool(
            controller.store.rename_flow, name, body.new_name
        )
    except FlowNotFound:
        raise HTTPException(status_code=404, detail=f"No flow named '{name}'")
    except FlowExists as exc:
        raise HTTPException(status_code=409, detail=str(exc))

    # The old name is nobody's namespace now; its values would sit there under a
    # flow that no longer exists. They repopulate under the new name on the next
    # run, so this is cleanup rather than a migration.
    await run_in_threadpool(controller.forget_flow, name)
    await controller.reload()
    return _detail(controller, renamed)


# -----------------------------------------------------------------------------
# Node source
# -----------------------------------------------------------------------------


@router.get("/{name}/nodes/{node_id}/source", response_model=NodeSource)
def read_node_source(
    name: str,
    node_id: str,
    controller: FlowControllerDep,
) -> Any:
    """Read a node's Python source, including unpublished edits."""
    definition = _read_flow(controller, name)
    ref = _source_ref(definition, node_id)
    if ref:
        try:
            return NodeSource(code=controller.store.read_lib_source(ref))
        except LibNotFound:
            raise HTTPException(status_code=404, detail=f"No shared node '{ref}'")
    return NodeSource(code=controller.store.read_node_source(name, node_id, draft=True))


@router.put("/{name}/nodes/{node_id}/source", response_model=NodeStatusPublic)
async def save_node_source(
    name: str,
    node_id: str,
    source: NodeSource,
    controller: FlowControllerDep,
) -> Any:
    """Save a node's source as an unpublished edit and report whether it loads.

    The answer comes from compiling the code rather than from the running
    pipeline: a draft is not deployed, and compiling is both faster and more
    precise about what the author just typed.

    A shared node writes to the library, so the fix reaches every flow using
    it — and that one is live immediately rather than waiting for a publish,
    because the copy is not any single flow's to hold back.
    """
    definition = _read_flow(controller, name)
    ref = _source_ref(definition, node_id)
    if ref:
        changed = await run_in_threadpool(
            controller.store.write_lib_source, ref, source.code
        )
        if changed:
            await controller.reload()
    else:
        await run_in_threadpool(
            controller.store.write_node_source, name, node_id, source.code, True
        )
    node_def = next((n for n in definition.nodes if n.id == node_id), None)
    device = (
        node_def.device
        if node_def is not None and node_def.device_policy == "require"
        else None
    )
    error = await run_in_threadpool(
        controller.compile_check, name, node_id, source.code, device
    )
    return NodeStatusPublic(
        id=f"{name}.{node_id}",
        status="error" if error else "active",
        error=error,
    )


@router.post("/{name}/nodes/{node_id}/share", response_model=FlowDetail)
async def share_node(
    name: str,
    node_id: str,
    body: ShareRequest,
    controller: FlowControllerDep,
) -> Any:
    """Move this node's code into the library so other flows can run it too."""
    _read_flow(controller, name)
    if not NAME_PATTERN.match(body.lib_name):
        raise HTTPException(
            status_code=400,
            detail=(
                "Use lowercase letters, digits and underscores, starting with a letter"
            ),
        )
    try:
        await run_in_threadpool(
            controller.store.share_node, name, node_id, body.lib_name
        )
    except LibExists as exc:
        raise HTTPException(status_code=409, detail=str(exc))
    return _detail(controller, _read_flow(controller, name))


@router.post("/{name}/nodes/{node_id}/unshare", response_model=FlowDetail)
async def unshare_node(
    name: str,
    node_id: str,
    controller: FlowControllerDep,
) -> Any:
    """Take a private copy of the shared code back into this flow."""
    _read_flow(controller, name)
    try:
        await run_in_threadpool(controller.store.unshare_node, name, node_id)
    except LibNotFound as exc:
        raise HTTPException(status_code=404, detail=str(exc))
    return _detail(controller, _read_flow(controller, name))


# -----------------------------------------------------------------------------
# Running, stopped, paused
# -----------------------------------------------------------------------------


@router.post("/{name}/start", response_model=FlowDetail)
async def start_flow(
    name: str, controller: FlowControllerDep, user: CurrentUser
) -> Any:
    """Let the engine run this flow again."""
    _read_flow(controller, name)
    await controller.set_enabled(name, True)
    _audit("started", name, user)
    return _detail(controller, _read_flow(controller, name))


@router.post("/{name}/stop", response_model=FlowDetail)
async def stop_flow(name: str, controller: FlowControllerDep, user: CurrentUser) -> Any:
    """Take this flow off the engine: no subscriptions, schedules or webhooks."""
    _read_flow(controller, name)
    await controller.set_enabled(name, False)
    _audit("stopped", name, user)
    return _detail(controller, _read_flow(controller, name))


@router.post("/{name}/pause", response_model=Message)
def pause_flow(name: str, controller: FlowControllerDep) -> Any:
    """Hold the flow's nodes so its messages can be stepped through."""
    _read_flow(controller, name)
    controller.pause_flow(name)
    return Message(message=f"Paused flow '{name}'")


@router.post("/{name}/resume", response_model=Message)
async def resume_flow(name: str, controller: FlowControllerDep) -> Any:
    """Let the flow carry on, running whatever was held back."""
    _read_flow(controller, name)
    await run_in_threadpool(controller.resume_flow, name)
    return Message(message=f"Resumed flow '{name}'")


@router.post("/{name}/step", response_model=Message)
async def step_flow(name: str, controller: FlowControllerDep) -> Any:
    """Run one message a pause is holding back, leaving the flow paused."""
    _read_flow(controller, name)
    node = await run_in_threadpool(controller.step_flow, name)
    if node is None:
        return Message(message=f"Nothing held back in flow '{name}'")
    return Message(message=f"Stepped '{node}'")


# -----------------------------------------------------------------------------
# Validation and execution
# -----------------------------------------------------------------------------


@router.post("/{name}/validate", response_model=ValidationResult)
def validate_flow(name: str, controller: FlowControllerDep) -> Any:
    """Report what would keep this flow from running."""
    _read_flow(controller, name)
    if controller.store.has_draft(name):
        return ValidationResult(issues=controller.preview(name).issues)
    return ValidationResult(issues=controller.flow_issues(name))


@router.post("/{name}/run", response_model=FlowStatePublic)
async def run_flow(
    name: str,
    body: RunRequest,
    controller: FlowControllerDep,
) -> Any:
    """Run every node of a flow once.

    With unpublished changes this runs the draft, so the button matches what is
    on the canvas. Nothing is deployed by running it.
    """
    _read_flow(controller, name)
    _require_enabled(controller, name)
    inputs = {qualify(name, key): value for key, value in body.inputs.items()}
    if controller.store.has_draft(name):
        await run_in_threadpool(controller.run_preview, name, inputs)
    else:
        await run_in_threadpool(controller.run_flow, name, inputs)
    return _flow_state(controller, name)


@router.post("/{name}/nodes/{node_id}/trigger", response_model=FlowStatePublic)
async def trigger_node(
    name: str,
    node_id: str,
    body: TriggerRequest,
    controller: FlowControllerDep,
) -> Any:
    """Feed values into a single node."""
    _require_enabled(controller, name)
    try:
        error = await run_in_threadpool(
            controller.trigger_node, f"{name}.{node_id}", body.values
        )
    except KeyError:
        raise HTTPException(
            status_code=404, detail=f"No node named '{node_id}' in flow '{name}'"
        )
    if error:
        # The canvas already has the failure from the bus; the person who
        # clicked gets to read it too, instead of a stack trace in the log.
        raise HTTPException(status_code=400, detail=error)
    return _flow_state(controller, name)


@router.post("/{name}/nodes/{node_id}/cancel", response_model=Message)
def cancel_node(name: str, node_id: str, controller: FlowControllerDep) -> Any:
    """Stop a node that is running right now, by killing the worker running it.

    Idempotent on purpose: by the time a click reaches here the node may well
    have finished, and that is the outcome that was asked for.
    """
    pool = controller.workers
    if pool is not None and pool.cancel(f"{name}.{node_id}"):
        return Message(message=f"Stopped '{node_id}'")
    return Message(message=f"'{node_id}' was not running")


@router.get("/{name}/state", response_model=FlowStatePublic)
def read_flow_state(name: str, controller: FlowControllerDep) -> Any:
    """The last value seen on every message of this flow."""
    _read_flow(controller, name)
    return _flow_state(controller, name)


@router.get("/{name}/history/{message}", response_model=MessageHistory)
def read_message_history(
    name: str,
    message: str,
    controller: FlowControllerDep,
) -> Any:
    """The recent values of one message, for plotting.

    ``message`` may be given bare or qualified; a message that never carried a
    number comes back with an empty series.
    """
    _read_flow(controller, name)
    key = qualify(name, message)
    return MessageHistory(
        message=key,
        numeric=as_number(controller.state.get(key)) is not None,
        points=[
            HistoryPoint(ts=ts, value=value)
            for ts, value in controller.state.history(key)
        ],
    )


# -----------------------------------------------------------------------------
# Live updates
# -----------------------------------------------------------------------------


@ws_router.websocket("/ws")
async def flow_events(websocket: WebSocket, token: str = "") -> None:
    """Stream values, node status and execution events as they happen.

    The token goes in the query string because browsers cannot set headers on
    a websocket handshake.
    """
    with Session(engine) as session:
        user = user_from_token(session, token)
    if user is None:
        await websocket.close(code=1008)
        return

    await websocket.accept()

    controller: FlowController | None = getattr(
        websocket.app.state, "flow_controller", None
    )
    if controller is not None:
        await websocket.send_json(
            {
                "type": "snapshot",
                "values": controller.values(),
                "nodes": [s.model_dump() for s in controller.node_statuses()],
                "issues": [i.model_dump() for i in controller.issues],
                "paused": controller.paused_flows(),
                "logs": list(event_bus.recent_logs),
            }
        )

    async with event_bus.subscribe() as queue:
        receiver = asyncio.create_task(websocket.receive_text())
        try:
            while True:
                sender = asyncio.create_task(queue.get())
                done, _ = await asyncio.wait(
                    {sender, receiver}, return_when=asyncio.FIRST_COMPLETED
                )
                if receiver in done:
                    # The client went away.
                    sender.cancel()
                    break
                await websocket.send_json(sender.result())
        except WebSocketDisconnect:
            pass
        finally:
            receiver.cancel()