# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: create_workflow_trigger_func.sql
import dataclasses

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


CREATE_WORKFLOW_TRIGGER = """-- name: create_workflow_trigger \\:exec

CREATE OR REPLACE FUNCTION trigger_workflow_run(
    p_workflow_run_id UUID,
    p_workflow_name TEXT,
    p_args JSON,
    p_queue TEXT,
    p_timeout_seconds INTEGER,
    p_idempotency_key TEXT
) RETURNS TABLE (
    workflow_run_id UUID,
    task_count INTEGER,
    success BOOLEAN,
    message TEXT
) AS $$
DECLARE
    v_workflow_record RECORD;
    v_node RECORD;
    v_task_id UUID;
    v_task_count INTEGER \\:= 0;
    v_workflow_dependencies UUID[];
    v_dep_durable_id TEXT;
    v_node_id_map JSONB \\:= '{}'\\:\\:JSONB;
BEGIN
    -- Validate timeout_seconds
    IF p_timeout_seconds IS NOT NULL AND p_timeout_seconds <= 0 THEN
        RETURN QUERY SELECT
            NULL\\:\\:UUID,
            0,
            FALSE,
            'Invalid timeout_seconds\\: must be NULL or greater than 0';
        RETURN;
    END IF;

    -- Start transaction
    BEGIN
        -- Get workflow definition including DAG structure
        SELECT * INTO v_workflow_record
        FROM hyrex_workflow
        WHERE workflow_name = p_workflow_name
        LIMIT 1;

        IF v_workflow_record IS NULL THEN
            RETURN QUERY SELECT
                NULL\\:\\:UUID,
                0,
                FALSE,
                'Workflow not found\\: ' || p_workflow_name;
            RETURN;
        END IF;

        -- Insert workflow run
        INSERT INTO hyrex_workflow_run (
            id,
            parent_id,
            workflow_name,
            args,
            queue,
            timeout_seconds,
            status,
            queued,
            last_heartbeat,
            idempotency_key
        )
        VALUES (
            p_workflow_run_id,
            NULL,
            p_workflow_name,
            p_args,
            p_queue,
            p_timeout_seconds,
            'RUNNING'\\:\\:workflow_run_status,
            NOW(),
            NOW(),
            p_idempotency_key
        );

        -- Create task runs from DAG structure
        IF v_workflow_record.dag_structure IS NOT NULL AND
           v_workflow_record.dag_structure->'nodes' IS NOT NULL THEN

            -- First pass\\: Create mapping of node IDs to task UUIDs
            FOR v_node IN
                SELECT value FROM jsonb_array_elements(v_workflow_record.dag_structure->'nodes')
            LOOP
                -- Generate a new UUID for this task
                v_task_id \\:= gen_random_uuid();

                -- Store the mapping from node ID to task UUID
                v_node_id_map \\:= v_node_id_map || jsonb_build_object(v_node.value->>'id', v_task_id\\:\\:TEXT);
            END LOOP;

            -- Second pass\\: Create task runs with proper dependencies
            FOR v_node IN
                SELECT value FROM jsonb_array_elements(v_workflow_record.dag_structure->'nodes')
            LOOP
                -- Get the task UUID we generated for this node
                v_task_id \\:= (v_node_id_map->>(v_node.value->>'id'))\\:\\:UUID;

                -- Build workflow_dependencies array based on edges
                v_workflow_dependencies \\:= ARRAY[]\\:\\:UUID[];

                -- Find all nodes that point to this node (dependencies)
                FOR v_dep_durable_id IN
                    SELECT DISTINCT value->>'from'
                    FROM jsonb_array_elements(v_workflow_record.dag_structure->'edges')
                    WHERE value->>'to' = v_node.value->>'id'
                LOOP
                    -- Map the dependency node ID to its task UUID
                    v_workflow_dependencies \\:= array_append(v_workflow_dependencies, (v_node_id_map->>v_dep_durable_id)\\:\\:UUID);
                END LOOP;

                -- Use create_task_run function with workflow parameters
                PERFORM create_task_run(
                    v_task_id,
                    v_task_id,  -- SDK uses same UUID for id and durable_id
                    v_task_id,  -- SDK uses same UUID for root_id too
                    NULL\\:\\:UUID,  -- No parent for workflow tasks
                    CASE
                        WHEN array_length(v_workflow_dependencies, 1) IS NULL OR array_length(v_workflow_dependencies, 1) = 0
                        THEN 'QUEUED'\\:\\:task_run_status
                        ELSE 'AWAIT_DEPS'\\:\\:task_run_status
                    END,
                    v_node.value->>'name',
                    p_args,  -- Pass workflow args to each task
                    p_queue,
                    3\\:\\:SMALLINT,  -- Default max retries
                    0\\:\\:SMALLINT,  -- Default priority
                    p_timeout_seconds,
                    NULL\\:\\:VARCHAR,  -- No idempotency key for individual tasks
                    NULL\\:\\:TIMESTAMP WITH TIME ZONE,  -- No scheduled_start for workflow tasks
                    p_workflow_run_id,
                    v_workflow_dependencies
                );

                v_task_count \\:= v_task_count + 1;
            END LOOP;
        END IF;

        -- Return success
        RETURN QUERY SELECT
            p_workflow_run_id,
            v_task_count,
            TRUE,
            'Workflow triggered successfully with ' || v_task_count || ' tasks';

    EXCEPTION WHEN OTHERS THEN
        -- Rollback will happen automatically
        RETURN QUERY SELECT
            NULL\\:\\:UUID,
            0,
            FALSE,
            'Error triggering workflow\\: ' || SQLERRM;
    END;
END;
$$ LANGUAGE plpgsql
"""


@dataclasses.dataclass()
class CreateWorkflowTriggerParams:

    pass

class Querier:
    def __init__(self, conn: sqlalchemy.engine.Connection):
        self._conn = conn

    def create_workflow_trigger(self, arg: CreateWorkflowTriggerParams) -> None:
        self._conn.execute(sqlalchemy.text(CREATE_WORKFLOW_TRIGGER))


class AsyncQuerier:
    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
        self._conn = conn

    async def create_workflow_trigger(self, arg: CreateWorkflowTriggerParams) -> None:
        await self._conn.execute(sqlalchemy.text(CREATE_WORKFLOW_TRIGGER))
