# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: create_task_run.sql
import dataclasses
import datetime
from typing import Any, List, Optional
import uuid

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


CREATE_TASK_RUN = """-- name: create_task_run \\:one
SELECT create_task_run(
    :p1\\:\\:UUID,
    :p2\\:\\:UUID,
    :p3\\:\\:UUID,
    :p4\\:\\:UUID,
    :p5\\:\\:task_run_status,
    :p6\\:\\:VARCHAR,
    :p7\\:\\:JSON,
    :p8\\:\\:VARCHAR,
    :p9\\:\\:SMALLINT,
    :p10\\:\\:SMALLINT,
    :p11\\:\\:INT,
    :p12\\:\\:VARCHAR,
    :p13\\:\\:TIMESTAMP WITH TIME ZONE,
    :p14\\:\\:UUID,
    :p15\\:\\:UUID[]
) AS id
"""


@dataclasses.dataclass()
class CreateTaskRunParams:
    id: uuid.UUID
    durable_id: uuid.UUID
    root_id: uuid.UUID
    parent_id: uuid.UUID
    status: Any
    task_name: str
    args: Any
    queue: str
    max_retries: int
    priority: int
    timeout_seconds: int
    idempotency_key: str
    scheduled_start: datetime.datetime
    workflow_run_id: uuid.UUID
    workflow_dependencies: List[uuid.UUID]


CREATE_TASK_RUN_FUNCTION = """-- name: create_task_run_function \\:exec
CREATE OR REPLACE FUNCTION create_task_run(
    p_id UUID,
    p_durable_id UUID,
    p_root_id UUID,
    p_parent_id UUID,
    p_status task_run_status,
    p_task_name VARCHAR,
    p_args JSON,
    p_queue VARCHAR,
    p_max_retries SMALLINT,
    p_priority SMALLINT,
    p_timeout_seconds INT DEFAULT NULL,
    p_idempotency_key VARCHAR DEFAULT NULL,
    p_scheduled_start TIMESTAMP WITH TIME ZONE DEFAULT NULL,
    p_workflow_run_id UUID DEFAULT NULL,
    p_workflow_dependencies UUID[] DEFAULT NULL
) RETURNS UUID AS $$
DECLARE
    v_task_id UUID;
    v_queued_timestamp TIMESTAMP WITH TIME ZONE;
BEGIN
    -- Validate inputs
    IF p_status = 'AWAIT_START_TIME' AND p_scheduled_start IS NULL THEN
        RAISE EXCEPTION 'scheduled_start is required when status is AWAIT_START_TIME';
    END IF;
    
    IF p_status NOT IN ('QUEUED', 'AWAIT_START_TIME', 'AWAIT_DEPS') THEN
        RAISE EXCEPTION 'Initial task status must be QUEUED, AWAIT_START_TIME, or AWAIT_DEPS';
    END IF;
    
    -- Set queued timestamp only for QUEUED status
    IF p_status = 'QUEUED' THEN
        v_queued_timestamp \\:= CURRENT_TIMESTAMP;
    ELSE
        v_queued_timestamp \\:= NULL;
    END IF;
    
    -- Insert the task
    INSERT INTO hyrex_task_run (
        id,
        durable_id,
        root_id,
        parent_id,
        status,
        task_name,
        args,
        queue,
        max_retries,
        priority,
        timeout_seconds,
        attempt_number,
        queued,
        idempotency_key,
        scheduled_start,
        workflow_run_id,
        workflow_dependencies
    )
    VALUES (
        p_id,
        p_durable_id,
        p_root_id,
        p_parent_id,
        p_status,
        p_task_name,
        p_args,
        p_queue,
        p_max_retries,
        p_priority,
        p_timeout_seconds,
        0,
        v_queued_timestamp,
        p_idempotency_key,
        p_scheduled_start,
        p_workflow_run_id,
        p_workflow_dependencies
    )
    ON CONFLICT (task_name, idempotency_key)
        WHERE idempotency_key IS NOT NULL
        DO NOTHING
    RETURNING id INTO v_task_id;
    
    -- Log idempotency conflict if it occurred
    IF v_task_id IS NULL AND p_idempotency_key IS NOT NULL THEN
        INSERT INTO hyrex_system_logs (id, timestamp, event_name, event_body)
        VALUES (
            gen_random_uuid(),
            CURRENT_TIMESTAMP,
            'task_idempotency_conflict',
            json_build_object(
                'task_name', p_task_name,
                'idempotency_key', p_idempotency_key,
                'attempted_task_id', p_id,
                'queue', p_queue,
                'status', p_status\\:\\:TEXT
            )
        );
    END IF;
    
    -- Return NULL if conflict occurred, otherwise return the new task ID
    RETURN v_task_id;
END;
$$ LANGUAGE plpgsql
"""


@dataclasses.dataclass()
class CreateTaskRunFunctionParams:

    pass

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

    def create_task_run(self, arg: CreateTaskRunParams) -> Optional[uuid.UUID]:
        row = self._conn.execute(sqlalchemy.text(CREATE_TASK_RUN), {
            "p1": arg.id,
            "p2": arg.durable_id,
            "p3": arg.root_id,
            "p4": arg.parent_id,
            "p5": arg.status,
            "p6": arg.task_name,
            "p7": arg.args,
            "p8": arg.queue,
            "p9": arg.max_retries,
            "p10": arg.priority,
            "p11": arg.timeout_seconds,
            "p12": arg.idempotency_key,
            "p13": arg.scheduled_start,
            "p14": arg.workflow_run_id,
            "p15": arg.workflow_dependencies,
        }).first()
        if row is None:
            return None
        return row[0]

    def create_task_run_function(self, arg: CreateTaskRunFunctionParams) -> None:
        self._conn.execute(sqlalchemy.text(CREATE_TASK_RUN_FUNCTION))


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

    async def create_task_run(self, arg: CreateTaskRunParams) -> Optional[uuid.UUID]:
        row = (await self._conn.execute(sqlalchemy.text(CREATE_TASK_RUN), {
            "p1": arg.id,
            "p2": arg.durable_id,
            "p3": arg.root_id,
            "p4": arg.parent_id,
            "p5": arg.status,
            "p6": arg.task_name,
            "p7": arg.args,
            "p8": arg.queue,
            "p9": arg.max_retries,
            "p10": arg.priority,
            "p11": arg.timeout_seconds,
            "p12": arg.idempotency_key,
            "p13": arg.scheduled_start,
            "p14": arg.workflow_run_id,
            "p15": arg.workflow_dependencies,
        })).first()
        if row is None:
            return None
        return row[0]

    async def create_task_run_function(self, arg: CreateTaskRunFunctionParams) -> None:
        await self._conn.execute(sqlalchemy.text(CREATE_TASK_RUN_FUNCTION))
