-- ============================================
-- Complete Async Task System in PostgreSQL
-- ============================================

-- 1. Create task queue and results tables
CREATE TABLE IF NOT EXISTS task_queue (
    id SERIAL PRIMARY KEY,
    task_type TEXT NOT NULL,
    task_data JSONB,
    status TEXT DEFAULT 'pending' CHECK (status IN ('pending', 'processing', 'completed', 'failed')),
    priority INTEGER DEFAULT 5,
    created_at TIMESTAMP DEFAULT NOW(),
    started_at TIMESTAMP,
    completed_at TIMESTAMP,
    result JSONB,
    error_message TEXT,
    retry_count INTEGER DEFAULT 0,
    max_retries INTEGER DEFAULT 3
);

CREATE INDEX idx_task_queue_status ON task_queue(status, priority DESC, created_at);

-- 2. Function to enqueue a task (returns immediately)
CREATE OR REPLACE FUNCTION enqueue_task(
    p_task_type TEXT,
    p_task_data JSONB DEFAULT '{}'::JSONB,
    p_priority INTEGER DEFAULT 5
)
RETURNS INTEGER AS $$
DECLARE
    v_task_id INTEGER;
BEGIN
    -- Insert task into queue
    INSERT INTO task_queue (task_type, task_data, status, priority)
    VALUES (p_task_type, p_task_data, 'pending', p_priority)
    RETURNING id INTO v_task_id;

    -- Notify workers that a new task is available
    PERFORM pg_notify('new_task', v_task_id::TEXT);

    RETURN v_task_id;
END;
$$ LANGUAGE plpgsql;

-- 3. Function to get next pending task (for workers)
CREATE OR REPLACE FUNCTION get_next_task()
RETURNS TABLE (
    task_id INTEGER,
    task_type TEXT,
    task_data JSONB
) AS $$
BEGIN
    RETURN QUERY
    UPDATE task_queue
    SET status = 'processing',
        started_at = NOW()
    WHERE id = (
        SELECT id FROM task_queue
        WHERE status = 'pending'
        ORDER BY priority DESC, created_at ASC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
    )
    RETURNING id, task_queue.task_type, task_queue.task_data;
END;
$$ LANGUAGE plpgsql;

-- 4. Function to mark task as completed
CREATE OR REPLACE FUNCTION complete_task(
    p_task_id INTEGER,
    p_result JSONB DEFAULT NULL
)
RETURNS VOID AS $$
BEGIN
    UPDATE task_queue
    SET status = 'completed',
        completed_at = NOW(),
        result = p_result
    WHERE id = p_task_id;

    -- Notify listeners that task is complete
    PERFORM pg_notify('task_completed', p_task_id::TEXT);
END;
$$ LANGUAGE plpgsql;

-- 5. Function to mark task as failed
CREATE OR REPLACE FUNCTION fail_task(
    p_task_id INTEGER,
    p_error_message TEXT
)
RETURNS VOID AS $$
DECLARE
    v_retry_count INTEGER;
    v_max_retries INTEGER;
BEGIN
    -- Get current retry info
    SELECT retry_count, max_retries
    INTO v_retry_count, v_max_retries
    FROM task_queue
    WHERE id = p_task_id;

    -- Check if we should retry
    IF v_retry_count < v_max_retries THEN
        UPDATE task_queue
        SET status = 'pending',
            retry_count = retry_count + 1,
            error_message = p_error_message
        WHERE id = p_task_id;
    ELSE
        UPDATE task_queue
        SET status = 'failed',
            completed_at = NOW(),
            error_message = p_error_message
        WHERE id = p_task_id;
    END IF;

    PERFORM pg_notify('task_failed', p_task_id::TEXT);
END;
$$ LANGUAGE plpgsql;

-- 6. Function to get task status
CREATE OR REPLACE FUNCTION get_task_status(p_task_id INTEGER)
RETURNS TABLE (
    id INTEGER,
    task_type TEXT,
    status TEXT,
    result JSONB,
    error_message TEXT,
    created_at TIMESTAMP,
    completed_at TIMESTAMP
) AS $$
BEGIN
    RETURN QUERY
    SELECT
        t.id,
        t.task_type,
        t.status,
        t.result,
        t.error_message,
        t.created_at,
        t.completed_at
    FROM task_queue t
    WHERE t.id = p_task_id;
END;
$$ LANGUAGE plpgsql;

-- 7. Example long-running function
CREATE OR REPLACE FUNCTION process_large_dataset(p_dataset_id INTEGER)
RETURNS JSONB AS $$
DECLARE
    v_result JSONB;
    v_count INTEGER;
BEGIN
    -- Simulate long-running processing
    PERFORM pg_sleep(5); -- 5 seconds

    -- Do actual work here
    -- SELECT COUNT(*) INTO v_count FROM large_table WHERE dataset_id = p_dataset_id;

    v_result := jsonb_build_object(
        'dataset_id', p_dataset_id,
        'processed_records', 1000,
        'status', 'success'
    );

    RETURN v_result;
END;
$$ LANGUAGE plpgsql;

-- 8. Cleanup old completed tasks
CREATE OR REPLACE FUNCTION cleanup_old_tasks(p_days_old INTEGER DEFAULT 7)
RETURNS INTEGER AS $$
DECLARE
    v_deleted INTEGER;
BEGIN
    DELETE FROM task_queue
    WHERE status IN ('completed', 'failed')
    AND completed_at < NOW() - INTERVAL '1 day' * p_days_old;

    GET DIAGNOSTICS v_deleted = ROW_COUNT;
    RETURN v_deleted;
END;
$$ LANGUAGE plpgsql;

-- ============================================
-- USAGE EXAMPLES
-- ============================================

-- Enqueue a task (returns immediately with task_id)
SELECT enqueue_task(
    'process_large_dataset',
    '{"dataset_id": 123, "options": {"verbose": true}}'::JSONB,
    10  -- high priority
);
-- Returns: task_id (e.g., 1)

-- Check task status
SELECT * FROM get_task_status(1);

-- Poll for completion
SELECT
    id,
    status,
    result,
    EXTRACT(EPOCH FROM (completed_at - created_at)) as duration_seconds
FROM task_queue
WHERE id = 1;

-- Get all pending tasks
SELECT * FROM task_queue WHERE status = 'pending' ORDER BY priority DESC;

-- Get all completed tasks from last hour
SELECT * FROM task_queue
WHERE status = 'completed'
AND completed_at > NOW() - INTERVAL '1 hour';

-- Clean up old tasks
SELECT cleanup_old_tasks(7); -- Remove tasks older than 7 days