durable/

ctx.rs

use std::pin::Pin;

use durable_db::entity::task::{
    ActiveModel as TaskActiveModel, Column as TaskColumn, Entity as Task,
};
use sea_orm::{
    ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseConnection, DatabaseTransaction,
    DbBackend, EntityTrait, QueryFilter, Set, Statement, TransactionTrait,
};
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::sync::atomic::{AtomicI32, Ordering};
use std::time::Duration;
use uuid::Uuid;

use crate::error::DurableError;

// ── Retry constants ───────────────────────────────────────────────────────────

const MAX_CHECKPOINT_RETRIES: u32 = 3;
const CHECKPOINT_RETRY_BASE_MS: u64 = 100;

/// Retry a fallible DB write with exponential backoff.
///
/// Calls `f` once immediately. On failure, retries up to `MAX_CHECKPOINT_RETRIES`
/// times with exponential backoff (100ms, 200ms, 400ms). Returns the FIRST
/// error if all retries are exhausted.
async fn retry_db_write<F, Fut>(mut f: F) -> Result<(), DurableError>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<(), DurableError>>,
{
    match f().await {
        Ok(()) => Ok(()),
        Err(first_err) => {
            for i in 0..MAX_CHECKPOINT_RETRIES {
                tokio::time::sleep(Duration::from_millis(
                    CHECKPOINT_RETRY_BASE_MS * 2u64.pow(i),
                ))
                .await;
                if f().await.is_ok() {
                    tracing::warn!(retry = i + 1, "checkpoint write succeeded on retry");
                    return Ok(());
                }
            }
            Err(first_err)
        }
    }
}

/// Policy for retrying a step on failure.
pub struct RetryPolicy {
    pub max_retries: u32,
    pub initial_backoff: std::time::Duration,
    pub backoff_multiplier: f64,
}

impl RetryPolicy {
    /// No retries — fails immediately on error.
    pub fn none() -> Self {
        Self {
            max_retries: 0,
            initial_backoff: std::time::Duration::from_secs(0),
            backoff_multiplier: 1.0,
        }
    }

    /// Exponential backoff: backoff doubles each retry.
    pub fn exponential(max_retries: u32, initial_backoff: std::time::Duration) -> Self {
        Self {
            max_retries,
            initial_backoff,
            backoff_multiplier: 2.0,
        }
    }

    /// Fixed backoff: same duration between all retries.
    pub fn fixed(max_retries: u32, backoff: std::time::Duration) -> Self {
        Self {
            max_retries,
            initial_backoff: backoff,
            backoff_multiplier: 1.0,
        }
    }
}

/// Context threaded through every workflow and step.
///
/// Users never create or manage task IDs. The SDK handles everything
/// via `(parent_id, name)` lookups — the unique constraint in the schema
/// guarantees exactly-once step creation.
pub struct Ctx {
    db: DatabaseConnection,
    task_id: Uuid,
    sequence: AtomicI32,
}

impl Ctx {
    // ── Workflow lifecycle (user-facing) ──────────────────────────

    /// Start or resume a root workflow by name.
    ///
    /// ```ignore
    /// let ctx = Ctx::start(&db, "ingest", json!({"crawl": "CC-2026"})).await?;
    /// ```
    pub async fn start(
        db: &DatabaseConnection,
        name: &str,
        input: Option<serde_json::Value>,
    ) -> Result<Self, DurableError> {
        let txn = db.begin().await?;
        // Workflows use plain find-or-create without FOR UPDATE SKIP LOCKED.
        // Multiple workers resuming the same workflow is fine — they share the
        // same task row and each step will be individually locked.
        let (task_id, _saved) =
            find_or_create_task(&txn, None, None, name, "WORKFLOW", input, false, None).await?;
        retry_db_write(|| set_status(&txn, task_id, "RUNNING")).await?;
        txn.commit().await?;
        Ok(Self {
            db: db.clone(),
            task_id,
            sequence: AtomicI32::new(0),
        })
    }

    /// Run a step. If already completed, returns saved output. Otherwise executes
    /// the closure, saves the result, and returns it.
    ///
    /// This method uses no retries (max_retries=0). For retries, use `step_with_retry`.
    ///
    /// ```ignore
    /// let count: u32 = ctx.step("fetch_count", || async { api.get_count().await }).await?;
    /// ```
    pub async fn step<T, F, Fut>(&self, name: &str, f: F) -> Result<T, DurableError>
    where
        T: Serialize + DeserializeOwned,
        F: FnOnce() -> Fut,
        Fut: std::future::Future<Output = Result<T, DurableError>>,
    {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);

        // Check if workflow is paused or cancelled before executing
        check_status(&self.db, self.task_id).await?;

        // Check if parent task's deadline has passed before executing
        check_deadline(&self.db, self.task_id).await?;

        // Begin transaction — the FOR UPDATE lock is held throughout step execution
        let txn = self.db.begin().await?;

        // Find or create — idempotent via UNIQUE(parent_id, name)
        // Returns (step_id, Option<saved_output>) where saved_output is Some iff COMPLETED.
        // Use FOR UPDATE SKIP LOCKED so only one worker can execute a given step.
        // max_retries=0: step() does not retry; use step_with_retry() for retries.
        let (step_id, saved_output) = find_or_create_task(
            &txn,
            Some(self.task_id),
            Some(seq),
            name,
            "STEP",
            None,
            true,
            Some(0),
        )
        .await?;

        // If already completed, replay from saved output
        if let Some(output) = saved_output {
            txn.commit().await?;
            let val: T = serde_json::from_value(output)?;
            tracing::debug!(step = name, seq, "replaying saved output");
            return Ok(val);
        }

        // Execute the step closure within the transaction (row lock is held)
        retry_db_write(|| set_status(&txn, step_id, "RUNNING")).await?;
        match f().await {
            Ok(val) => {
                let json = serde_json::to_value(&val)?;
                retry_db_write(|| complete_task(&txn, step_id, json.clone())).await?;
                txn.commit().await?;
                tracing::debug!(step = name, seq, "step completed");
                Ok(val)
            }
            Err(e) => {
                let err_msg = e.to_string();
                retry_db_write(|| fail_task(&txn, step_id, &err_msg)).await?;
                txn.commit().await?;
                Err(e)
            }
        }
    }

    /// Run a DB-only step inside a single Postgres transaction.
    ///
    /// Both the user's DB work and the checkpoint save happen in the same transaction,
    /// ensuring atomicity. If the closure returns an error, both the user writes and
    /// the checkpoint are rolled back.
    ///
    /// ```ignore
    /// let count: u32 = ctx.transaction("upsert_batch", |tx| Box::pin(async move {
    ///     do_db_work(tx).await
    /// })).await?;
    /// ```
    pub async fn transaction<T, F>(&self, name: &str, f: F) -> Result<T, DurableError>
    where
        T: Serialize + DeserializeOwned + Send,
        F: for<'tx> FnOnce(
                &'tx DatabaseTransaction,
            ) -> Pin<
                Box<dyn std::future::Future<Output = Result<T, DurableError>> + Send + 'tx>,
            > + Send,
    {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);

        // Check if workflow is paused or cancelled before executing
        check_status(&self.db, self.task_id).await?;

        // Find or create the step task record OUTSIDE the transaction.
        // This is idempotent (UNIQUE constraint) and must exist before we begin.
        let (step_id, saved_output) = find_or_create_task(
            &self.db,
            Some(self.task_id),
            Some(seq),
            name,
            "TRANSACTION",
            None,
            false,
            None,
        )
        .await?;

        // If already completed, replay from saved output.
        if let Some(output) = saved_output {
            let val: T = serde_json::from_value(output)?;
            tracing::debug!(step = name, seq, "replaying saved transaction output");
            return Ok(val);
        }

        // Begin the transaction — set_status, user work, and complete_task all happen atomically.
        let tx = self.db.begin().await?;

        set_status(&tx, step_id, "RUNNING").await?;

        match f(&tx).await {
            Ok(val) => {
                let json = serde_json::to_value(&val)?;
                complete_task(&tx, step_id, json).await?;
                tx.commit().await?;
                tracing::debug!(step = name, seq, "transaction step committed");
                Ok(val)
            }
            Err(e) => {
                // Rollback happens automatically when tx is dropped.
                // Record the failure on the main connection (outside the rolled-back tx).
                drop(tx);
                fail_task(&self.db, step_id, &e.to_string()).await?;
                Err(e)
            }
        }
    }

    /// Start or resume a child workflow. Returns a new `Ctx` scoped to the child.
    ///
    /// ```ignore
    /// let child_ctx = ctx.child("embed_batch", json!({"vectors": 1000})).await?;
    /// // use child_ctx.step(...) for steps inside the child
    /// child_ctx.complete(json!({"done": true})).await?;
    /// ```
    pub async fn child(
        &self,
        name: &str,
        input: Option<serde_json::Value>,
    ) -> Result<Self, DurableError> {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);

        // Check if workflow is paused or cancelled before executing
        check_status(&self.db, self.task_id).await?;

        let txn = self.db.begin().await?;
        // Child workflows also use plain find-or-create without locking.
        let (child_id, _saved) = find_or_create_task(
            &txn,
            Some(self.task_id),
            Some(seq),
            name,
            "WORKFLOW",
            input,
            false,
            None,
        )
        .await?;

        // If child already completed, return a Ctx that will replay
        // (the caller should check is_completed() or just run steps which will replay)
        retry_db_write(|| set_status(&txn, child_id, "RUNNING")).await?;
        txn.commit().await?;

        Ok(Self {
            db: self.db.clone(),
            task_id: child_id,
            sequence: AtomicI32::new(0),
        })
    }

    /// Check if this workflow/child was already completed (for skipping in parent).
    pub async fn is_completed(&self) -> Result<bool, DurableError> {
        let status = get_status(&self.db, self.task_id).await?;
        Ok(status.as_deref() == Some("COMPLETED"))
    }

    /// Get the saved output if this task is completed.
    pub async fn get_output<T: DeserializeOwned>(&self) -> Result<Option<T>, DurableError> {
        match get_output(&self.db, self.task_id).await? {
            Some(val) => Ok(Some(serde_json::from_value(val)?)),
            None => Ok(None),
        }
    }

    /// Mark this workflow as completed with an output value.
    pub async fn complete<T: Serialize>(&self, output: &T) -> Result<(), DurableError> {
        let json = serde_json::to_value(output)?;
        let db = &self.db;
        let task_id = self.task_id;
        retry_db_write(|| complete_task(db, task_id, json.clone())).await
    }

    /// Run a step with a configurable retry policy.
    ///
    /// Unlike `step()`, the closure must implement `Fn` (not `FnOnce`) since
    /// it may be called multiple times on retry. Retries happen in-process with
    /// configurable backoff between attempts.
    ///
    /// ```ignore
    /// let result: u32 = ctx
    ///     .step_with_retry("call_api", RetryPolicy::exponential(3, Duration::from_secs(1)), || async {
    ///         api.call().await
    ///     })
    ///     .await?;
    /// ```
    pub async fn step_with_retry<T, F, Fut>(
        &self,
        name: &str,
        policy: RetryPolicy,
        f: F,
    ) -> Result<T, DurableError>
    where
        T: Serialize + DeserializeOwned,
        F: Fn() -> Fut,
        Fut: std::future::Future<Output = Result<T, DurableError>>,
    {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);

        // Check if workflow is paused or cancelled before executing
        check_status(&self.db, self.task_id).await?;

        // Find or create — idempotent via UNIQUE(parent_id, name)
        // Set max_retries from policy when creating.
        // No locking here — retry logic handles re-execution in-process.
        let (step_id, saved_output) = find_or_create_task(
            &self.db,
            Some(self.task_id),
            Some(seq),
            name,
            "STEP",
            None,
            false,
            Some(policy.max_retries),
        )
        .await?;

        // If already completed, replay from saved output
        if let Some(output) = saved_output {
            let val: T = serde_json::from_value(output)?;
            tracing::debug!(step = name, seq, "replaying saved output");
            return Ok(val);
        }

        // Get current retry state from DB (for resume across restarts)
        let (mut retry_count, max_retries) = get_retry_info(&self.db, step_id).await?;

        // Retry loop
        loop {
            // Re-check status before each retry attempt
            check_status(&self.db, self.task_id).await?;
            set_status(&self.db, step_id, "RUNNING").await?;
            match f().await {
                Ok(val) => {
                    let json = serde_json::to_value(&val)?;
                    complete_task(&self.db, step_id, json).await?;
                    tracing::debug!(step = name, seq, retry_count, "step completed");
                    return Ok(val);
                }
                Err(e) => {
                    if retry_count < max_retries {
                        // Increment retry count and reset to PENDING
                        retry_count = increment_retry_count(&self.db, step_id).await?;
                        tracing::debug!(
                            step = name,
                            seq,
                            retry_count,
                            max_retries,
                            "step failed, retrying"
                        );

                        // Compute backoff duration
                        let backoff = if policy.initial_backoff.is_zero() {
                            std::time::Duration::ZERO
                        } else {
                            let factor = policy
                                .backoff_multiplier
                                .powi((retry_count - 1) as i32)
                                .max(1.0);
                            let millis =
                                (policy.initial_backoff.as_millis() as f64 * factor) as u64;
                            std::time::Duration::from_millis(millis)
                        };

                        if !backoff.is_zero() {
                            tokio::time::sleep(backoff).await;
                        }
                    } else {
                        // Exhausted retries — mark FAILED
                        fail_task(&self.db, step_id, &e.to_string()).await?;
                        tracing::debug!(
                            step = name,
                            seq,
                            retry_count,
                            "step exhausted retries, marked FAILED"
                        );
                        return Err(e);
                    }
                }
            }
        }
    }

    /// Mark this workflow as failed.
    pub async fn fail(&self, error: &str) -> Result<(), DurableError> {
        let db = &self.db;
        let task_id = self.task_id;
        retry_db_write(|| fail_task(db, task_id, error)).await
    }

    /// Set the timeout for this task in milliseconds.
    ///
    /// If a task stays RUNNING longer than timeout_ms, it will be eligible
    /// for recovery by `Executor::recover()`.
    ///
    /// If the task is currently RUNNING and has a `started_at`, this also
    /// computes and stores `deadline_epoch_ms = started_at_epoch_ms + timeout_ms`.
    pub async fn set_timeout(&self, timeout_ms: i64) -> Result<(), DurableError> {
        let sql = format!(
            "UPDATE durable.task \
             SET timeout_ms = {timeout_ms}, \
                 deadline_epoch_ms = CASE \
                     WHEN status = 'RUNNING' AND started_at IS NOT NULL \
                     THEN EXTRACT(EPOCH FROM started_at) * 1000 + {timeout_ms} \
                     ELSE deadline_epoch_ms \
                 END \
             WHERE id = '{}'",
            self.task_id
        );
        self.db
            .execute(Statement::from_string(DbBackend::Postgres, sql))
            .await?;
        Ok(())
    }

    /// Start or resume a root workflow with a timeout in milliseconds.
    ///
    /// Equivalent to calling `Ctx::start()` followed by `ctx.set_timeout(timeout_ms)`.
    pub async fn start_with_timeout(
        db: &DatabaseConnection,
        name: &str,
        input: Option<serde_json::Value>,
        timeout_ms: i64,
    ) -> Result<Self, DurableError> {
        let ctx = Self::start(db, name, input).await?;
        ctx.set_timeout(timeout_ms).await?;
        Ok(ctx)
    }

    // ── Workflow control (management API) ─────────────────────────

    /// Pause a workflow by ID. Sets status to PAUSED and recursively
    /// cascades to all PENDING/RUNNING descendants (children, grandchildren, etc.).
    ///
    /// Only workflows in PENDING or RUNNING status can be paused.
    pub async fn pause(db: &DatabaseConnection, task_id: Uuid) -> Result<(), DurableError> {
        let model = Task::find_by_id(task_id).one(db).await?;
        let model =
            model.ok_or_else(|| DurableError::custom(format!("task {task_id} not found")))?;

        match model.status.as_str() {
            "PENDING" | "RUNNING" => {}
            status => {
                return Err(DurableError::custom(format!(
                    "cannot pause task in {status} status"
                )));
            }
        }

        // Pause the task itself and all descendants in one recursive CTE
        let sql = format!(
            "WITH RECURSIVE descendants AS ( \
                 SELECT id FROM durable.task WHERE id = '{task_id}' \
                 UNION ALL \
                 SELECT t.id FROM durable.task t \
                 INNER JOIN descendants d ON t.parent_id = d.id \
             ) \
             UPDATE durable.task SET status = 'PAUSED' \
             WHERE id IN (SELECT id FROM descendants) \
               AND status IN ('PENDING', 'RUNNING')"
        );
        db.execute(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

        tracing::info!(%task_id, "workflow paused");
        Ok(())
    }

    /// Resume a paused workflow by ID. Sets status back to RUNNING and
    /// recursively cascades to all PAUSED descendants (resetting them to PENDING).
    pub async fn resume(db: &DatabaseConnection, task_id: Uuid) -> Result<(), DurableError> {
        let model = Task::find_by_id(task_id).one(db).await?;
        let model =
            model.ok_or_else(|| DurableError::custom(format!("task {task_id} not found")))?;

        if model.status != "PAUSED" {
            return Err(DurableError::custom(format!(
                "cannot resume task in {} status (must be PAUSED)",
                model.status
            )));
        }

        // Resume the root task back to RUNNING
        let sql = format!(
            "UPDATE durable.task SET status = 'RUNNING' WHERE id = '{task_id}'"
        );
        db.execute(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

        // Recursively resume all PAUSED descendants back to PENDING
        let cascade_sql = format!(
            "WITH RECURSIVE descendants AS ( \
                 SELECT id FROM durable.task WHERE parent_id = '{task_id}' \
                 UNION ALL \
                 SELECT t.id FROM durable.task t \
                 INNER JOIN descendants d ON t.parent_id = d.id \
             ) \
             UPDATE durable.task SET status = 'PENDING' \
             WHERE id IN (SELECT id FROM descendants) \
               AND status = 'PAUSED'"
        );
        db.execute(Statement::from_string(DbBackend::Postgres, cascade_sql))
            .await?;

        tracing::info!(%task_id, "workflow resumed");
        Ok(())
    }

    /// Cancel a workflow by ID. Sets status to CANCELLED and recursively
    /// cascades to all non-terminal descendants.
    ///
    /// Cancellation is terminal — a cancelled workflow cannot be resumed.
    pub async fn cancel(db: &DatabaseConnection, task_id: Uuid) -> Result<(), DurableError> {
        let model = Task::find_by_id(task_id).one(db).await?;
        let model =
            model.ok_or_else(|| DurableError::custom(format!("task {task_id} not found")))?;

        match model.status.as_str() {
            "COMPLETED" | "FAILED" | "CANCELLED" => {
                return Err(DurableError::custom(format!(
                    "cannot cancel task in {} status",
                    model.status
                )));
            }
            _ => {}
        }

        // Cancel the task itself and all non-terminal descendants in one recursive CTE
        let sql = format!(
            "WITH RECURSIVE descendants AS ( \
                 SELECT id FROM durable.task WHERE id = '{task_id}' \
                 UNION ALL \
                 SELECT t.id FROM durable.task t \
                 INNER JOIN descendants d ON t.parent_id = d.id \
             ) \
             UPDATE durable.task SET status = 'CANCELLED', completed_at = now() \
             WHERE id IN (SELECT id FROM descendants) \
               AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED')"
        );
        db.execute(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

        tracing::info!(%task_id, "workflow cancelled");
        Ok(())
    }

    // ── Accessors ────────────────────────────────────────────────

    pub fn db(&self) -> &DatabaseConnection {
        &self.db
    }

    pub fn task_id(&self) -> Uuid {
        self.task_id
    }

    pub fn next_sequence(&self) -> i32 {
        self.sequence.fetch_add(1, Ordering::SeqCst)
    }
}

// ── Internal SQL helpers ─────────────────────────────────────────────

/// Find an existing task by (parent_id, name) or create a new one.
///
/// Returns `(task_id, Option<saved_output>)`:
/// - `saved_output` is `Some(json)` when the task is COMPLETED (replay path).
/// - `saved_output` is `None` when the task is new or in-progress.
///
/// When `lock` is `true` and an existing non-completed task is found, this
/// function attempts to acquire a `FOR UPDATE SKIP LOCKED` row lock. If
/// another worker holds the lock, `DurableError::StepLocked` is returned so
/// the caller can skip execution rather than double-firing side effects.
///
/// When `lock` is `false`, a plain SELECT is used (appropriate for workflow
/// and child-workflow creation where concurrent start is safe).
///
/// When `lock` is `true`, the caller MUST call this within a transaction so
/// the row lock is held throughout step execution.
///
/// `max_retries`: if Some, overrides the schema default when creating a new task.
#[allow(clippy::too_many_arguments)]
async fn find_or_create_task(
    db: &impl ConnectionTrait,
    parent_id: Option<Uuid>,
    sequence: Option<i32>,
    name: &str,
    kind: &str,
    input: Option<serde_json::Value>,
    lock: bool,
    max_retries: Option<u32>,
) -> Result<(Uuid, Option<serde_json::Value>), DurableError> {
    let parent_eq = match parent_id {
        Some(p) => format!("= '{p}'"),
        None => "IS NULL".to_string(),
    };
    let parent_sql = match parent_id {
        Some(p) => format!("'{p}'"),
        None => "NULL".to_string(),
    };

    if lock {
        // Locking path (for steps): we need exactly-once execution.
        //
        // Strategy:
        // 1. INSERT the row with ON CONFLICT DO NOTHING — idempotent creation.
        //    If another transaction is concurrently inserting the same row,
        //    Postgres will block here until that transaction commits or rolls
        //    back, ensuring we never see a phantom "not found" for a row being
        //    inserted.
        // 2. Attempt FOR UPDATE SKIP LOCKED — if we just inserted the row we
        //    should get it back; if the row existed and is locked by another
        //    worker we get nothing.
        // 3. If SKIP LOCKED returns empty, return StepLocked.

        let new_id = Uuid::new_v4();
        let seq_sql = match sequence {
            Some(s) => s.to_string(),
            None => "NULL".to_string(),
        };
        let input_sql = match &input {
            Some(v) => format!("'{}'", serde_json::to_string(v)?),
            None => "NULL".to_string(),
        };

        let max_retries_sql = match max_retries {
            Some(r) => r.to_string(),
            None => "3".to_string(), // schema default
        };

        // Step 1: insert-or-skip
        let insert_sql = format!(
            "INSERT INTO durable.task (id, parent_id, sequence, name, kind, status, input, max_retries) \
             VALUES ('{new_id}', {parent_sql}, {seq_sql}, '{name}', '{kind}', 'PENDING', {input_sql}, {max_retries_sql}) \
             ON CONFLICT (parent_id, name) DO NOTHING"
        );
        db.execute(Statement::from_string(DbBackend::Postgres, insert_sql))
            .await?;

        // Step 2: lock the row (ours or pre-existing)
        let lock_sql = format!(
            "SELECT id, status, output FROM durable.task \
             WHERE parent_id {parent_eq} AND name = '{name}' \
             FOR UPDATE SKIP LOCKED"
        );
        let row = db
            .query_one(Statement::from_string(DbBackend::Postgres, lock_sql))
            .await?;

        if let Some(row) = row {
            let id: Uuid = row
                .try_get_by_index(0)
                .map_err(|e| DurableError::custom(e.to_string()))?;
            let status: String = row
                .try_get_by_index(1)
                .map_err(|e| DurableError::custom(e.to_string()))?;
            let output: Option<serde_json::Value> = row.try_get_by_index(2).ok();

            if status == "COMPLETED" {
                // Replay path — return saved output
                return Ok((id, output));
            }
            // Task exists and we hold the lock — proceed to execute
            return Ok((id, None));
        }

        // Step 3: SKIP LOCKED returned empty — another worker holds the lock
        Err(DurableError::StepLocked(name.to_string()))
    } else {
        // Plain find without locking — safe for workflow-level operations.
        // Multiple workers resuming the same workflow is fine; individual
        // steps will be locked when executed.
        let mut query = Task::find().filter(TaskColumn::Name.eq(name));
        query = match parent_id {
            Some(p) => query.filter(TaskColumn::ParentId.eq(p)),
            None => query.filter(TaskColumn::ParentId.is_null()),
        };
        let existing = query.one(db).await?;

        if let Some(model) = existing {
            if model.status == "COMPLETED" {
                return Ok((model.id, model.output));
            }
            return Ok((model.id, None));
        }

        // Task does not exist — create it
        let id = Uuid::new_v4();
        let new_task = TaskActiveModel {
            id: Set(id),
            parent_id: Set(parent_id),
            sequence: Set(sequence),
            name: Set(name.to_string()),
            kind: Set(kind.to_string()),
            status: Set("PENDING".to_string()),
            input: Set(input),
            max_retries: Set(max_retries.map(|r| r as i32).unwrap_or(3)),
            ..Default::default()
        };
        new_task.insert(db).await?;

        Ok((id, None))
    }
}

async fn get_output(
    db: &impl ConnectionTrait,
    task_id: Uuid,
) -> Result<Option<serde_json::Value>, DurableError> {
    let model = Task::find_by_id(task_id)
        .filter(TaskColumn::Status.eq("COMPLETED"))
        .one(db)
        .await?;

    Ok(model.and_then(|m| m.output))
}

async fn get_status(
    db: &impl ConnectionTrait,
    task_id: Uuid,
) -> Result<Option<String>, DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    Ok(model.map(|m| m.status))
}

/// Returns (retry_count, max_retries) for a task.
async fn get_retry_info(
    db: &DatabaseConnection,
    task_id: Uuid,
) -> Result<(u32, u32), DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    match model {
        Some(m) => Ok((m.retry_count as u32, m.max_retries as u32)),
        None => Err(DurableError::custom(format!(
            "task {task_id} not found when reading retry info"
        ))),
    }
}

/// Increment retry_count and reset status to PENDING. Returns the new retry_count.
async fn increment_retry_count(
    db: &DatabaseConnection,
    task_id: Uuid,
) -> Result<u32, DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    match model {
        Some(m) => {
            let new_count = m.retry_count + 1;
            let mut active: TaskActiveModel = m.into();
            active.retry_count = Set(new_count);
            active.status = Set("PENDING".to_string());
            active.error = Set(None);
            active.completed_at = Set(None);
            active.update(db).await?;
            Ok(new_count as u32)
        }
        None => Err(DurableError::custom(format!(
            "task {task_id} not found when incrementing retry count"
        ))),
    }
}

// NOTE: set_status uses raw SQL because SeaORM cannot express the CASE expression
// for conditional deadline_epoch_ms computation or COALESCE on started_at.
async fn set_status(
    db: &impl ConnectionTrait,
    task_id: Uuid,
    status: &str,
) -> Result<(), DurableError> {
    let sql = format!(
        "UPDATE durable.task \
         SET status = '{status}', \
             started_at = COALESCE(started_at, now()), \
             deadline_epoch_ms = CASE \
                 WHEN '{status}' = 'RUNNING' AND timeout_ms IS NOT NULL \
                 THEN EXTRACT(EPOCH FROM now()) * 1000 + timeout_ms \
                 ELSE deadline_epoch_ms \
             END \
         WHERE id = '{task_id}'"
    );
    db.execute(Statement::from_string(DbBackend::Postgres, sql))
        .await?;
    Ok(())
}

/// Check if the task is paused or cancelled. Returns an error if so.
async fn check_status(db: &DatabaseConnection, task_id: Uuid) -> Result<(), DurableError> {
    let status = get_status(db, task_id).await?;
    match status.as_deref() {
        Some("PAUSED") => Err(DurableError::Paused(format!("task {task_id} is paused"))),
        Some("CANCELLED") => Err(DurableError::Cancelled(format!("task {task_id} is cancelled"))),
        _ => Ok(()),
    }
}

/// Check if the parent task's deadline has passed. Returns `DurableError::Timeout` if so.
async fn check_deadline(db: &DatabaseConnection, task_id: Uuid) -> Result<(), DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    if let Some(m) = model
        && let Some(deadline_ms) = m.deadline_epoch_ms
    {
        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_millis() as i64)
            .unwrap_or(0);
        if now_ms > deadline_ms {
            return Err(DurableError::Timeout("task deadline exceeded".to_string()));
        }
    }

    Ok(())
}

async fn complete_task(
    db: &impl ConnectionTrait,
    task_id: Uuid,
    output: serde_json::Value,
) -> Result<(), DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    if let Some(m) = model {
        let mut active: TaskActiveModel = m.into();
        active.status = Set("COMPLETED".to_string());
        active.output = Set(Some(output));
        active.completed_at = Set(Some(chrono::Utc::now().into()));
        active.update(db).await?;
    }
    Ok(())
}

async fn fail_task(
    db: &impl ConnectionTrait,
    task_id: Uuid,
    error: &str,
) -> Result<(), DurableError> {
    let model = Task::find_by_id(task_id).one(db).await?;

    if let Some(m) = model {
        let mut active: TaskActiveModel = m.into();
        active.status = Set("FAILED".to_string());
        active.error = Set(Some(error.to_string()));
        active.completed_at = Set(Some(chrono::Utc::now().into()));
        active.update(db).await?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU32, Ordering};

    /// test_retry_db_write_succeeds_first_try: a closure that always succeeds
    /// should be called exactly once and return Ok.
    #[tokio::test]
    async fn test_retry_db_write_succeeds_first_try() {
        let call_count = Arc::new(AtomicU32::new(0));
        let cc = call_count.clone();
        let result = retry_db_write(|| {
            let c = cc.clone();
            async move {
                c.fetch_add(1, Ordering::SeqCst);
                Ok::<(), DurableError>(())
            }
        })
        .await;
        assert!(result.is_ok());
        assert_eq!(call_count.load(Ordering::SeqCst), 1);
    }

    /// test_retry_db_write_succeeds_after_transient_failure: a closure that
    /// fails twice then succeeds should return Ok and be called 3 times.
    #[tokio::test]
    async fn test_retry_db_write_succeeds_after_transient_failure() {
        let call_count = Arc::new(AtomicU32::new(0));
        let cc = call_count.clone();
        let result = retry_db_write(|| {
            let c = cc.clone();
            async move {
                let n = c.fetch_add(1, Ordering::SeqCst);
                if n < 2 {
                    Err(DurableError::Db(sea_orm::DbErr::Custom(
                        "transient".to_string(),
                    )))
                } else {
                    Ok(())
                }
            }
        })
        .await;
        assert!(result.is_ok());
        assert_eq!(call_count.load(Ordering::SeqCst), 3);
    }

    /// test_retry_db_write_exhausts_retries: a closure that always fails should
    /// be called 1 + MAX_CHECKPOINT_RETRIES times total then return an error.
    #[tokio::test]
    async fn test_retry_db_write_exhausts_retries() {
        let call_count = Arc::new(AtomicU32::new(0));
        let cc = call_count.clone();
        let result = retry_db_write(|| {
            let c = cc.clone();
            async move {
                c.fetch_add(1, Ordering::SeqCst);
                Err::<(), DurableError>(DurableError::Db(sea_orm::DbErr::Custom(
                    "always fails".to_string(),
                )))
            }
        })
        .await;
        assert!(result.is_err());
        // 1 initial attempt + MAX_CHECKPOINT_RETRIES retry attempts
        assert_eq!(
            call_count.load(Ordering::SeqCst),
            1 + MAX_CHECKPOINT_RETRIES
        );
    }

    /// test_retry_db_write_returns_original_error: when all retries are
    /// exhausted the FIRST error is returned, not the last retry error.
    #[tokio::test]
    async fn test_retry_db_write_returns_original_error() {
        let call_count = Arc::new(AtomicU32::new(0));
        let cc = call_count.clone();
        let result = retry_db_write(|| {
            let c = cc.clone();
            async move {
                let n = c.fetch_add(1, Ordering::SeqCst);
                Err::<(), DurableError>(DurableError::Db(sea_orm::DbErr::Custom(format!(
                    "error-{}",
                    n
                ))))
            }
        })
        .await;
        let err = result.unwrap_err();
        // The message of the first error contains "error-0"
        assert!(
            err.to_string().contains("error-0"),
            "expected first error (error-0), got: {err}"
        );
    }
}