DocsGPT/application/api/user/workflows/routes.py

"""Workflow management routes."""

from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Set

from flask import current_app, request
from flask_restx import Namespace, Resource

from application.api.user.base import (
    workflow_edges_collection,
    workflow_nodes_collection,
    workflows_collection,
)
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
from application.storage.db.repositories.workflows import WorkflowsRepository
from application.core.json_schema_utils import (
    JsonSchemaValidationError,
    normalize_json_schema_payload,
)
from application.core.model_utils import get_model_capabilities
from application.api.user.utils import (
    check_resource_ownership,
    error_response,
    get_user_id,
    require_auth,
    require_fields,
    safe_db_operation,
    success_response,
    validate_object_id,
)

workflows_ns = Namespace("workflows", path="/api")


def _workflow_error_response(message: str, err: Exception):
    current_app.logger.error(f"{message}: {err}", exc_info=True)
    return error_response(message)


# ---------------------------------------------------------------------------
# Postgres dual-write helpers
#
# Workflows are unusual relative to other Phase 3 tables: a single user
# action (create / update) writes to three collections in concert
# (workflows + workflow_nodes + workflow_edges) and the edges reference
# nodes by user-provided string ids. The Postgres mirror needs to:
#
# 1. Run all three writes inside one PG transaction (so the just-created
#    nodes are visible when we resolve their UUIDs for the edge insert).
# 2. Translate edge source_id/target_id strings → workflow_nodes.id UUIDs
#    after the bulk_create returns them.
#
# Each helper opens exactly one ``dual_write`` call (one PG txn) and uses
# the connection from whichever repo it was instantiated with to spin up
# any sibling repos it needs.
# ---------------------------------------------------------------------------


def _dual_write_workflow_create(
    mongo_workflow_id: str,
    user_id: str,
    name: str,
    description: str,
    nodes_data: List[Dict],
    edges_data: List[Dict],
    graph_version: int = 1,
) -> None:
    """Mirror a Mongo workflow create into Postgres."""

    def _do(repo: WorkflowsRepository) -> None:
        conn = repo._conn
        wf = repo.create(
            user_id,
            name,
            description=description,
            legacy_mongo_id=mongo_workflow_id,
        )
        _write_graph(conn, wf["id"], graph_version, nodes_data, edges_data)

    dual_write(WorkflowsRepository, _do)


def _dual_write_workflow_update(
    mongo_workflow_id: str,
    user_id: str,
    name: str,
    description: str,
    nodes_data: List[Dict],
    edges_data: List[Dict],
    next_graph_version: int,
) -> None:
    """Mirror a Mongo workflow update into Postgres.

    Mirrors the Mongo route: insert the new graph_version's nodes/edges,
    bump the workflow's name/description/current_graph_version, then drop
    every other graph_version's nodes/edges.
    """

    def _do(repo: WorkflowsRepository) -> None:
        conn = repo._conn
        wf = _resolve_pg_workflow(conn, mongo_workflow_id)
        if wf is None:
            return
        _write_graph(conn, wf["id"], next_graph_version, nodes_data, edges_data)
        repo.update(wf["id"], user_id, {
            "name": name,
            "description": description,
            "current_graph_version": next_graph_version,
        })
        WorkflowNodesRepository(conn).delete_other_versions(
            wf["id"], next_graph_version,
        )
        WorkflowEdgesRepository(conn).delete_other_versions(
            wf["id"], next_graph_version,
        )

    dual_write(WorkflowsRepository, _do)


def _dual_write_workflow_delete(mongo_workflow_id: str, user_id: str) -> None:
    """Mirror a Mongo workflow delete into Postgres.

    The CASCADE on workflows.id → workflow_nodes/workflow_edges takes
    care of the children automatically.
    """

    def _do(repo: WorkflowsRepository) -> None:
        wf = _resolve_pg_workflow(repo._conn, mongo_workflow_id)
        if wf is not None:
            repo.delete(wf["id"], user_id)

    dual_write(WorkflowsRepository, _do)


def _resolve_pg_workflow(conn, mongo_workflow_id: str) -> Optional[Dict]:
    """Look up a Postgres workflow by its Mongo ObjectId string."""
    from sqlalchemy import text as _text
    row = conn.execute(
        _text("SELECT id FROM workflows WHERE legacy_mongo_id = :legacy_id"),
        {"legacy_id": mongo_workflow_id},
    ).fetchone()
    return {"id": str(row[0])} if row else None


def _write_graph(
    conn,
    pg_workflow_id: str,
    graph_version: int,
    nodes_data: List[Dict],
    edges_data: List[Dict],
) -> None:
    """Bulk-create nodes + edges for one graph version inside one txn.

    Edges arrive with source/target as user-provided node-id strings
    (the same shape the Mongo route stores). We bulk-insert nodes first,
    capture their ``node_id → UUID`` map from the returned rows, then
    translate edge source/target strings to those UUIDs before the edge
    bulk insert. Edges referencing missing nodes are dropped (logged).
    """
    nodes_repo = WorkflowNodesRepository(conn)
    edges_repo = WorkflowEdgesRepository(conn)

    if nodes_data:
        created_nodes = nodes_repo.bulk_create(
            pg_workflow_id, graph_version,
            [
                {
                    "node_id": n["id"],
                    "node_type": n["type"],
                    "title": n.get("title", ""),
                    "description": n.get("description", ""),
                    "position": n.get("position", {"x": 0, "y": 0}),
                    "config": n.get("data", {}),
                    "legacy_mongo_id": n.get("legacy_mongo_id"),
                }
                for n in nodes_data
            ],
        )
        node_uuid_by_str = {n["node_id"]: n["id"] for n in created_nodes}
    else:
        node_uuid_by_str = {}

    if edges_data:
        translated_edges: List[Dict] = []
        for e in edges_data:
            src = e.get("source")
            tgt = e.get("target")
            from_uuid = node_uuid_by_str.get(src)
            to_uuid = node_uuid_by_str.get(tgt)
            if not from_uuid or not to_uuid:
                current_app.logger.warning(
                    "PG dual-write: dropping edge %s; node refs unresolved "
                    "(source=%s, target=%s)",
                    e.get("id"), src, tgt,
                )
                continue
            translated_edges.append({
                "edge_id": e["id"],
                "from_node_id": from_uuid,
                "to_node_id": to_uuid,
                "source_handle": e.get("sourceHandle"),
                "target_handle": e.get("targetHandle"),
            })
        if translated_edges:
            edges_repo.bulk_create(pg_workflow_id, graph_version, translated_edges)


def serialize_workflow(w: Dict) -> Dict:
    """Serialize workflow document to API response format."""
    return {
        "id": str(w["_id"]),
        "name": w.get("name"),
        "description": w.get("description"),
        "created_at": w["created_at"].isoformat() if w.get("created_at") else None,
        "updated_at": w["updated_at"].isoformat() if w.get("updated_at") else None,
    }


def serialize_node(n: Dict) -> Dict:
    """Serialize workflow node document to API response format."""
    return {
        "id": n["id"],
        "type": n["type"],
        "title": n.get("title"),
        "description": n.get("description"),
        "position": n.get("position"),
        "data": n.get("config", {}),
    }


def serialize_edge(e: Dict) -> Dict:
    """Serialize workflow edge document to API response format."""
    return {
        "id": e["id"],
        "source": e.get("source_id"),
        "target": e.get("target_id"),
        "sourceHandle": e.get("source_handle"),
        "targetHandle": e.get("target_handle"),
    }


def get_workflow_graph_version(workflow: Dict) -> int:
    """Get current graph version with legacy fallback."""
    raw_version = workflow.get("current_graph_version", 1)
    try:
        version = int(raw_version)
        return version if version > 0 else 1
    except (ValueError, TypeError):
        return 1


def fetch_graph_documents(collection, workflow_id: str, graph_version: int) -> List[Dict]:
    """Fetch graph docs for active version, with fallback for legacy unversioned data."""
    docs = list(
        collection.find({"workflow_id": workflow_id, "graph_version": graph_version})
    )
    if docs:
        return docs
    if graph_version == 1:
        return list(
            collection.find(
                {"workflow_id": workflow_id, "graph_version": {"$exists": False}}
            )
        )
    return docs


def validate_json_schema_payload(
    json_schema: Any,
) -> tuple[Optional[Dict[str, Any]], Optional[str]]:
    """Validate and normalize optional JSON schema payload for structured output."""
    if json_schema is None:
        return None, None
    try:
        return normalize_json_schema_payload(json_schema), None
    except JsonSchemaValidationError as exc:
        return None, str(exc)


def normalize_agent_node_json_schemas(nodes: List[Dict]) -> List[Dict]:
    """Normalize agent-node JSON schema payloads before persistence."""
    normalized_nodes: List[Dict] = []
    for node in nodes:
        if not isinstance(node, dict):
            normalized_nodes.append(node)
            continue

        normalized_node = dict(node)
        if normalized_node.get("type") != "agent":
            normalized_nodes.append(normalized_node)
            continue

        raw_config = normalized_node.get("data")
        if not isinstance(raw_config, dict) or "json_schema" not in raw_config:
            normalized_nodes.append(normalized_node)
            continue

        normalized_config = dict(raw_config)
        try:
            normalized_config["json_schema"] = normalize_json_schema_payload(
                raw_config.get("json_schema")
            )
        except JsonSchemaValidationError:
            # Validation runs before normalization; keep original on unexpected shape.
            normalized_config["json_schema"] = raw_config.get("json_schema")
        normalized_node["data"] = normalized_config
        normalized_nodes.append(normalized_node)

    return normalized_nodes


def validate_workflow_structure(nodes: List[Dict], edges: List[Dict]) -> List[str]:
    """Validate workflow graph structure."""
    errors = []

    if not nodes:
        errors.append("Workflow must have at least one node")
        return errors

    start_nodes = [n for n in nodes if n.get("type") == "start"]
    if len(start_nodes) != 1:
        errors.append("Workflow must have exactly one start node")

    end_nodes = [n for n in nodes if n.get("type") == "end"]
    if not end_nodes:
        errors.append("Workflow must have at least one end node")

    node_ids = {n.get("id") for n in nodes}
    node_map = {n.get("id"): n for n in nodes}
    end_ids = {n.get("id") for n in end_nodes}

    for edge in edges:
        source_id = edge.get("source")
        target_id = edge.get("target")
        if source_id not in node_ids:
            errors.append(f"Edge references non-existent source: {source_id}")
        if target_id not in node_ids:
            errors.append(f"Edge references non-existent target: {target_id}")

    if start_nodes:
        start_id = start_nodes[0].get("id")
        if not any(e.get("source") == start_id for e in edges):
            errors.append("Start node must have at least one outgoing edge")

    condition_nodes = [n for n in nodes if n.get("type") == "condition"]
    for cnode in condition_nodes:
        cnode_id = cnode.get("id")
        cnode_title = cnode.get("title", cnode_id)
        outgoing = [e for e in edges if e.get("source") == cnode_id]
        if len(outgoing) < 2:
            errors.append(
                f"Condition node '{cnode_title}' must have at least 2 outgoing edges"
            )
        node_data = cnode.get("data", {}) or {}
        cases = node_data.get("cases", [])
        if not isinstance(cases, list):
            cases = []
        if not cases or not any(
            isinstance(c, dict) and str(c.get("expression", "")).strip() for c in cases
        ):
            errors.append(
                f"Condition node '{cnode_title}' must have at least one case with an expression"
            )

        case_handles: Set[str] = set()
        duplicate_case_handles: Set[str] = set()
        for case in cases:
            if not isinstance(case, dict):
                continue
            raw_handle = case.get("sourceHandle", "")
            handle = raw_handle.strip() if isinstance(raw_handle, str) else ""
            if not handle:
                errors.append(
                    f"Condition node '{cnode_title}' has a case without a branch handle"
                )
                continue
            if handle in case_handles:
                duplicate_case_handles.add(handle)
            case_handles.add(handle)

        for handle in duplicate_case_handles:
            errors.append(
                f"Condition node '{cnode_title}' has duplicate case handle '{handle}'"
            )

        outgoing_by_handle: Dict[str, List[Dict]] = {}
        for out_edge in outgoing:
            raw_handle = out_edge.get("sourceHandle", "")
            handle = raw_handle.strip() if isinstance(raw_handle, str) else ""
            outgoing_by_handle.setdefault(handle, []).append(out_edge)

        for handle, handle_edges in outgoing_by_handle.items():
            if not handle:
                errors.append(
                    f"Condition node '{cnode_title}' has an outgoing edge without sourceHandle"
                )
                continue
            if handle != "else" and handle not in case_handles:
                errors.append(
                    f"Condition node '{cnode_title}' has a connection from unknown branch '{handle}'"
                )
            if len(handle_edges) > 1:
                errors.append(
                    f"Condition node '{cnode_title}' has multiple outgoing edges from branch '{handle}'"
                )

        if "else" not in outgoing_by_handle:
            errors.append(f"Condition node '{cnode_title}' must have an 'else' branch")

        for case in cases:
            if not isinstance(case, dict):
                continue
            raw_handle = case.get("sourceHandle", "")
            handle = raw_handle.strip() if isinstance(raw_handle, str) else ""
            if not handle:
                continue

            raw_expression = case.get("expression", "")
            has_expression = isinstance(raw_expression, str) and bool(
                raw_expression.strip()
            )
            has_outgoing = bool(outgoing_by_handle.get(handle))
            if has_expression and not has_outgoing:
                errors.append(
                    f"Condition node '{cnode_title}' case '{handle}' has an expression but no outgoing edge"
                )
            if not has_expression and has_outgoing:
                errors.append(
                    f"Condition node '{cnode_title}' case '{handle}' has an outgoing edge but no expression"
                )

        for handle, handle_edges in outgoing_by_handle.items():
            if not handle:
                continue
            for out_edge in handle_edges:
                target = out_edge.get("target")
                if target and not _can_reach_end(target, edges, node_map, end_ids):
                    errors.append(
                        f"Branch '{handle}' of condition '{cnode_title}' "
                        f"must eventually reach an end node"
                    )

    agent_nodes = [n for n in nodes if n.get("type") == "agent"]
    for agent_node in agent_nodes:
        agent_title = agent_node.get("title", agent_node.get("id", "unknown"))
        raw_config = agent_node.get("data", {}) or {}
        if not isinstance(raw_config, dict):
            errors.append(f"Agent node '{agent_title}' has invalid configuration")
            continue
        normalized_schema, schema_error = validate_json_schema_payload(
            raw_config.get("json_schema")
        )
        has_json_schema = normalized_schema is not None

        model_id = raw_config.get("model_id")
        if has_json_schema and isinstance(model_id, str) and model_id.strip():
            capabilities = get_model_capabilities(model_id.strip())
            if capabilities and not capabilities.get("supports_structured_output", False):
                errors.append(
                    f"Agent node '{agent_title}' selected model does not support structured output"
                )
        if schema_error:
            errors.append(f"Agent node '{agent_title}' JSON schema {schema_error}")

    for node in nodes:
        if not node.get("id"):
            errors.append("All nodes must have an id")
        if not node.get("type"):
            errors.append(f"Node {node.get('id', 'unknown')} must have a type")

    return errors


def _can_reach_end(
    node_id: str, edges: List[Dict], node_map: Dict, end_ids: set, visited: set = None
) -> bool:
    if visited is None:
        visited = set()
    if node_id in end_ids:
        return True
    if node_id in visited or node_id not in node_map:
        return False
    visited.add(node_id)
    outgoing = [e.get("target") for e in edges if e.get("source") == node_id]
    return any(_can_reach_end(t, edges, node_map, end_ids, visited) for t in outgoing if t)


def create_workflow_nodes(
    workflow_id: str, nodes_data: List[Dict], graph_version: int
) -> List[Dict]:
    """Insert workflow nodes into Mongo and return rows with Mongo ids."""
    if nodes_data:
        mongo_nodes = [
            {
                "id": n["id"],
                "workflow_id": workflow_id,
                "graph_version": graph_version,
                "type": n["type"],
                "title": n.get("title", ""),
                "description": n.get("description", ""),
                "position": n.get("position", {"x": 0, "y": 0}),
                "config": n.get("data", {}),
            }
            for n in nodes_data
        ]
        result = workflow_nodes_collection.insert_many(mongo_nodes)
        return [
            {**node, "legacy_mongo_id": str(inserted_id)}
            for node, inserted_id in zip(nodes_data, result.inserted_ids)
        ]
    return []


def create_workflow_edges(
    workflow_id: str, edges_data: List[Dict], graph_version: int
) -> None:
    """Insert workflow edges into database."""
    if edges_data:
        workflow_edges_collection.insert_many(
            [
                {
                    "id": e["id"],
                    "workflow_id": workflow_id,
                    "graph_version": graph_version,
                    "source_id": e.get("source"),
                    "target_id": e.get("target"),
                    "source_handle": e.get("sourceHandle"),
                    "target_handle": e.get("targetHandle"),
                }
                for e in edges_data
            ]
        )


@workflows_ns.route("/workflows")
class WorkflowList(Resource):

    @require_auth
    @require_fields(["name"])
    def post(self):
        """Create a new workflow with nodes and edges."""
        user_id = get_user_id()
        data = request.get_json()

        name = data.get("name", "").strip()
        nodes_data = data.get("nodes", [])
        edges_data = data.get("edges", [])

        validation_errors = validate_workflow_structure(nodes_data, edges_data)
        if validation_errors:
            return error_response(
                "Workflow validation failed", errors=validation_errors
            )
        nodes_data = normalize_agent_node_json_schemas(nodes_data)

        now = datetime.now(timezone.utc)
        workflow_doc = {
            "name": name,
            "description": data.get("description", ""),
            "user": user_id,
            "created_at": now,
            "updated_at": now,
            "current_graph_version": 1,
        }

        result, error = safe_db_operation(
            lambda: workflows_collection.insert_one(workflow_doc),
            "Failed to create workflow",
        )
        if error:
            return error

        workflow_id = str(result.inserted_id)

        try:
            created_nodes = create_workflow_nodes(workflow_id, nodes_data, 1)
            create_workflow_edges(workflow_id, edges_data, 1)
        except Exception as err:
            workflow_nodes_collection.delete_many({"workflow_id": workflow_id})
            workflow_edges_collection.delete_many({"workflow_id": workflow_id})
            workflows_collection.delete_one({"_id": result.inserted_id})
            return _workflow_error_response("Failed to create workflow structure", err)

        _dual_write_workflow_create(
            workflow_id,
            user_id,
            name,
            data.get("description", ""),
            created_nodes,
            edges_data,
        )

        return success_response({"id": workflow_id}, 201)


@workflows_ns.route("/workflows/<string:workflow_id>")
class WorkflowDetail(Resource):

    @require_auth
    def get(self, workflow_id: str):
        """Get workflow details with nodes and edges."""
        user_id = get_user_id()
        obj_id, error = validate_object_id(workflow_id, "Workflow")
        if error:
            return error

        workflow, error = check_resource_ownership(
            workflows_collection, obj_id, user_id, "Workflow"
        )
        if error:
            return error

        graph_version = get_workflow_graph_version(workflow)
        nodes = fetch_graph_documents(
            workflow_nodes_collection, workflow_id, graph_version
        )
        edges = fetch_graph_documents(
            workflow_edges_collection, workflow_id, graph_version
        )

        return success_response(
            {
                "workflow": serialize_workflow(workflow),
                "nodes": [serialize_node(n) for n in nodes],
                "edges": [serialize_edge(e) for e in edges],
            }
        )

    @require_auth
    @require_fields(["name"])
    def put(self, workflow_id: str):
        """Update workflow and replace nodes/edges."""
        user_id = get_user_id()
        obj_id, error = validate_object_id(workflow_id, "Workflow")
        if error:
            return error

        workflow, error = check_resource_ownership(
            workflows_collection, obj_id, user_id, "Workflow"
        )
        if error:
            return error

        data = request.get_json()
        name = data.get("name", "").strip()
        nodes_data = data.get("nodes", [])
        edges_data = data.get("edges", [])

        validation_errors = validate_workflow_structure(nodes_data, edges_data)
        if validation_errors:
            return error_response(
                "Workflow validation failed", errors=validation_errors
            )
        nodes_data = normalize_agent_node_json_schemas(nodes_data)

        current_graph_version = get_workflow_graph_version(workflow)
        next_graph_version = current_graph_version + 1
        try:
            created_nodes = create_workflow_nodes(
                workflow_id, nodes_data, next_graph_version,
            )
            create_workflow_edges(workflow_id, edges_data, next_graph_version)
        except Exception as err:
            workflow_nodes_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": next_graph_version}
            )
            workflow_edges_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": next_graph_version}
            )
            return _workflow_error_response("Failed to update workflow structure", err)

        now = datetime.now(timezone.utc)
        _, error = safe_db_operation(
            lambda: workflows_collection.update_one(
                {"_id": obj_id},
                {
                    "$set": {
                        "name": name,
                        "description": data.get("description", ""),
                        "updated_at": now,
                        "current_graph_version": next_graph_version,
                    }
                },
            ),
            "Failed to update workflow",
        )
        if error:
            workflow_nodes_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": next_graph_version}
            )
            workflow_edges_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": next_graph_version}
            )
            return error

        try:
            workflow_nodes_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": {"$ne": next_graph_version}}
            )
            workflow_edges_collection.delete_many(
                {"workflow_id": workflow_id, "graph_version": {"$ne": next_graph_version}}
            )
        except Exception as cleanup_err:
            current_app.logger.warning(
                f"Failed to clean old workflow graph versions for {workflow_id}: {cleanup_err}"
            )

        _dual_write_workflow_update(
            workflow_id,
            user_id,
            name,
            data.get("description", ""),
            created_nodes,
            edges_data,
            next_graph_version,
        )

        return success_response()

    @require_auth
    def delete(self, workflow_id: str):
        """Delete workflow and its graph."""
        user_id = get_user_id()
        obj_id, error = validate_object_id(workflow_id, "Workflow")
        if error:
            return error

        workflow, error = check_resource_ownership(
            workflows_collection, obj_id, user_id, "Workflow"
        )
        if error:
            return error

        try:
            workflow_nodes_collection.delete_many({"workflow_id": workflow_id})
            workflow_edges_collection.delete_many({"workflow_id": workflow_id})
            workflows_collection.delete_one({"_id": workflow["_id"], "user": user_id})
        except Exception as err:
            return _workflow_error_response("Failed to delete workflow", err)

        _dual_write_workflow_delete(workflow_id, user_id)

        return success_response()