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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


CREATE_TRANSITION_TASK_RUN_STATE_FUNC = """-- name: create_transition_task_run_state_func \\:exec
DO $$
BEGIN
    -- ──────────────────────────────────────────────────────────────
    --  is_allowed(from_state, to_state)
    --  • Returns TRUE iff the edge exists in the Hyrex task-run FSM
    --  • Immutable, so PostgreSQL inlines it into plans (O-(1) test)
    -- ──────────────────────────────────────────────────────────────
    CREATE OR REPLACE FUNCTION is_allowed(
        p_from task_run_status,
        p_to   task_run_status
    ) RETURNS boolean
        IMMUTABLE
        LANGUAGE sql AS $func$
    SELECT (p_from, p_to) IN (
        -- ──────────────────────
        -- AwaitDeps transitions
        -- ──────────────────────
        ROW('AWAIT_DEPS',       'AWAIT_START_TIME'),
        ROW('AWAIT_DEPS',       'CANCELED'),
        ROW('AWAIT_DEPS',       'QUEUED'),
        ROW('AWAIT_DEPS',       'SKIPPED'),

        -- ─────────────────────────────────
        -- AwaitStartTime transitions
        -- ─────────────────────────────────
        ROW('AWAIT_START_TIME', 'QUEUED'),
        ROW('AWAIT_START_TIME', 'CANCELED'),

        -- ────────────────
        -- Queued → …
        -- ────────────────
        ROW('QUEUED',           'RUNNING'),
        ROW('QUEUED',           'CANCELED'),

        -- ──────────────────────────────
        -- Running → terminal / cancel
        -- ──────────────────────────────
        ROW('RUNNING',          'STOPPED'),
        ROW('RUNNING',          'SUCCESS'),
        ROW('RUNNING',          'FAILED'),
        ROW('RUNNING',          'LOST'),
        ROW('RUNNING',          'UP_FOR_CANCEL'),

        -- ─────────────────────────
        -- UpForCancel → terminal
        -- ─────────────────────────
        ROW('UP_FOR_CANCEL',    'FAILED'),
        ROW('UP_FOR_CANCEL',    'CANCELED')
    );
    $func$;

    CREATE OR REPLACE FUNCTION transition_task_state(
        p_id        uuid,
        p_next      task_run_status           -- enum type
    ) RETURNS hyrex_task_run
        LANGUAGE plpgsql AS $func$
    DECLARE
        v_old  hyrex_task_run;        -- row before change
        v_new  hyrex_task_run;        -- row after  change
        v_error_msg text;
        v_error_detail text;
        v_error_hint text;
    BEGIN
        ------------------------------------------------------------------
        -- 1. Lock the row so concurrent transitions queue up.
        ------------------------------------------------------------------
        SELECT * INTO v_old
        FROM hyrex_task_run
        WHERE id = p_id
            FOR UPDATE;

        IF NOT FOUND THEN
            RAISE EXCEPTION 'task % not found', p_id;
        END IF;

        ------------------------------------------------------------------
        -- 2. Validate the transition against the FSM.
        ------------------------------------------------------------------
        IF NOT is_allowed(v_old.status, p_next) THEN      -- helper shown below
            RAISE EXCEPTION
                'illegal state change\\: % → %', v_old.status, p_next;
        END IF;

        ------------------------------------------------------------------
        -- 3. Apply the change + timestamp side-effects in one shot.
        ------------------------------------------------------------------
        UPDATE hyrex_task_run
        SET status   = p_next,
            queued   = CASE WHEN p_next = 'QUEUED'   THEN clock_timestamp() ELSE queued   END,
            started  = CASE WHEN p_next = 'RUNNING'  THEN clock_timestamp() ELSE started  END,
            finished = CASE WHEN p_next IN ('SUCCESS','FAILED',
                                            'LOST','STOPPED','SKIPPED','CANCELED')
                                THEN clock_timestamp()
                            ELSE finished
                END,
            last_heartbeat = clock_timestamp()
        WHERE id = p_id
        RETURNING * INTO v_new;

        RETURN v_new;

    EXCEPTION
        WHEN OTHERS THEN
            -- Capture error details
            GET STACKED DIAGNOSTICS 
                v_error_msg = MESSAGE_TEXT,
                v_error_detail = PG_EXCEPTION_DETAIL,
                v_error_hint = PG_EXCEPTION_HINT;
            
            -- Log the error to system logs
            INSERT INTO hyrex_system_logs (id, timestamp, event_name, event_body)
            VALUES (
                gen_random_uuid(),
                clock_timestamp(),
                'task_state_transition_failed',
                jsonb_build_object(
                    'task_id', p_id,
                    'from_state', v_old.status,
                    'to_state', p_next,
                    'error_message', v_error_msg,
                    'error_detail', v_error_detail,
                    'error_hint', v_error_hint,
                    'sqlstate', SQLSTATE
                )
            );
            
            -- Re-raise the original exception
            RAISE;
    END;
    $func$;
END $$
"""


@dataclasses.dataclass()
class CreateTransitionTaskRunStateFuncParams:

    pass

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

    def create_transition_task_run_state_func(self, arg: CreateTransitionTaskRunStateFuncParams) -> None:
        self._conn.execute(sqlalchemy.text(CREATE_TRANSITION_TASK_RUN_STATE_FUNC))


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

    async def create_transition_task_run_state_func(self, arg: CreateTransitionTaskRunStateFuncParams) -> None:
        await self._conn.execute(sqlalchemy.text(CREATE_TRANSITION_TASK_RUN_STATE_FUNC))
