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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


CREATE_CONDITIONALLY_RETRY_TASK_FUNC = """-- name: create_conditionally_retry_task_func \\:exec
CREATE OR REPLACE FUNCTION conditionally_retry_task(
    p_existing_task_id UUID,
    p_new_task_id UUID,
    p_timeout_seconds INT DEFAULT NULL
) RETURNS TABLE (
    success BOOLEAN,
    message TEXT,
    new_task_id UUID
) AS $$
DECLARE
    v_task_record RECORD;
    v_retry_created BOOLEAN \\:= FALSE;
BEGIN
    -- Get the existing task details
    SELECT * INTO v_task_record
    FROM hyrex_task_run
    WHERE id = p_existing_task_id
    FOR UPDATE;
    
    IF v_task_record IS NULL THEN
        RETURN QUERY SELECT FALSE, 'Task not found', NULL\\:\\:UUID;
        RETURN;
    END IF;
    
    -- Check if retry is allowed
    IF v_task_record.attempt_number >= v_task_record.max_retries THEN
        -- Log max retries reached
        INSERT INTO hyrex_system_logs (id, timestamp, event_name, event_body)
        VALUES (
            gen_random_uuid(),
            CURRENT_TIMESTAMP,
            'task_max_retries_reached',
            json_build_object(
                'task_id', p_existing_task_id,
                'task_name', v_task_record.task_name,
                'attempt_number', v_task_record.attempt_number,
                'max_retries', v_task_record.max_retries
            )
        );
        
        RETURN QUERY SELECT FALSE, 'Max retries reached', NULL\\:\\:UUID;
        RETURN;
    END IF;
    
    -- Create the retry task
    INSERT INTO hyrex_task_run (
        id,
        durable_id,
        root_id,
        parent_id,
        workflow_run_id,
        workflow_dependencies,
        task_name,
        args,
        queue,
        max_retries,
        priority,
        timeout_seconds,
        status,
        attempt_number,
        queued,
        idempotency_key,
        scheduled_start
    )
    VALUES (
        p_new_task_id,
        v_task_record.durable_id,
        v_task_record.root_id,
        v_task_record.parent_id,
        v_task_record.workflow_run_id,
        v_task_record.workflow_dependencies,
        v_task_record.task_name,
        v_task_record.args,
        v_task_record.queue,
        v_task_record.max_retries,
        v_task_record.priority,
        COALESCE(p_timeout_seconds, v_task_record.timeout_seconds),
        'QUEUED'\\:\\:task_run_status,
        v_task_record.attempt_number + 1,
        CURRENT_TIMESTAMP,
        NULL,  -- No idempotency key for retries
        NULL   -- No scheduled start for retries
    );
    
    -- Log the retry
    INSERT INTO hyrex_system_logs (id, timestamp, event_name, event_body)
    VALUES (
        gen_random_uuid(),
        CURRENT_TIMESTAMP,
        'task_retry_created',
        json_build_object(
            'original_task_id', p_existing_task_id,
            'new_task_id', p_new_task_id,
            'task_name', v_task_record.task_name,
            'attempt_number', v_task_record.attempt_number + 1,
            'max_retries', v_task_record.max_retries
        )
    );
    
    RETURN QUERY SELECT TRUE, 'Retry task created successfully', p_new_task_id;
END;
$$ LANGUAGE plpgsql
"""


@dataclasses.dataclass()
class CreateConditionallyRetryTaskFuncParams:

    pass

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

    def create_conditionally_retry_task_func(self, arg: CreateConditionallyRetryTaskFuncParams) -> None:
        self._conn.execute(sqlalchemy.text(CREATE_CONDITIONALLY_RETRY_TASK_FUNC))


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

    async def create_conditionally_retry_task_func(self, arg: CreateConditionallyRetryTaskFuncParams) -> None:
        await self._conn.execute(sqlalchemy.text(CREATE_CONDITIONALLY_RETRY_TASK_FUNC))
