cloacina 0.11.1

/*
 *  Copyright 2025-2026 Colliery Software
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

//! Recovery operations for orphaned and failed tasks.

use super::{RetryStats, TaskExecutionDAL};
use crate::database::schema::unified::task_executions;
use crate::database::universal_types::{UniversalTimestamp, UniversalUuid};
use crate::error::ValidationError;
use crate::models::task_execution::TaskExecution;
use diesel::prelude::*;

impl<'a> TaskExecutionDAL<'a> {
    // NOTE (CLOACI-T-0914): `get_orphaned_tasks` (an unfiltered scan of all
    // Running rows, claim-blind) was removed — it had no production callers
    // and misleadingly suggested a recovery path that doesn't exist. Orphan
    // recovery is the StaleClaimSweeper's job: heartbeat-based, claim-aware,
    // and capped via `reset_task_for_recovery` / `mark_abandoned` below.

    /// Resets a task from "Running" to "Ready" state for recovery.
    pub async fn reset_task_for_recovery(
        &self,
        task_id: UniversalUuid,
    ) -> Result<(), ValidationError> {
        let now = UniversalTimestamp::now();
        crate::interact_on_backend!(self.dal, |conn| {
            diesel::update(task_executions::table.find(task_id))
                .set((
                    task_executions::status.eq("Ready"),
                    task_executions::started_at.eq(None::<UniversalTimestamp>),
                    task_executions::recovery_attempts.eq(task_executions::recovery_attempts + 1),
                    task_executions::last_recovery_at.eq(Some(now)),
                    task_executions::updated_at.eq(now),
                ))
                .execute(conn)
        })?;

        Ok(())
    }

    /// Checks if a workflow should be marked as failed due to abandoned tasks.
    pub async fn check_workflow_failure(
        &self,
        workflow_execution_id: UniversalUuid,
    ) -> Result<bool, ValidationError> {
        let failed_count: i64 = crate::interact_on_backend!(self.dal, |conn| {
            task_executions::table
                .filter(task_executions::workflow_execution_id.eq(workflow_execution_id))
                .filter(task_executions::status.eq("Failed"))
                .filter(task_executions::error_details.like("ABANDONED:%"))
                .count()
                .get_result(conn)
        })?;

        Ok(failed_count > 0)
    }

    /// Calculates retry statistics for a specific workflow execution.
    pub async fn get_retry_stats(
        &self,
        workflow_execution_id: UniversalUuid,
    ) -> Result<RetryStats, ValidationError> {
        // This method is backend-agnostic since it processes data in memory
        let tasks = self
            .get_all_tasks_for_workflow(workflow_execution_id)
            .await?;

        let mut stats = RetryStats::default();

        for task in tasks {
            if task.attempt > 1 {
                stats.tasks_with_retries += 1;
                stats.total_retries += task.attempt - 1;
            }

            if task.attempt > stats.max_attempts_used {
                stats.max_attempts_used = task.attempt;
            }

            if task.status == "Failed" && task.attempt >= task.max_attempts {
                stats.tasks_exhausted_retries += 1;
            }
        }

        Ok(stats)
    }

    /// Retrieves tasks that have exceeded their retry limit.
    pub async fn get_exhausted_retry_tasks(
        &self,
        workflow_execution_id: UniversalUuid,
    ) -> Result<Vec<TaskExecution>, ValidationError> {
        // This method is backend-agnostic since it filters in memory
        let tasks = self
            .get_all_tasks_for_workflow(workflow_execution_id)
            .await?;

        let exhausted_tasks: Vec<TaskExecution> = tasks
            .into_iter()
            .filter(|task| task.status == "Failed" && task.attempt >= task.max_attempts)
            .collect();

        Ok(exhausted_tasks)
    }
}

#[cfg(test)]
mod tests {
    use crate::dal::DAL;
    use crate::database::universal_types::UniversalUuid;
    use crate::database::Database;
    use crate::models::task_execution::NewTaskExecution;
    use crate::models::workflow_execution::NewWorkflowExecution;

    #[cfg(feature = "sqlite")]
    async fn unique_dal() -> DAL {
        let url = format!(
            "file:recovery_test_{}?mode=memory&cache=shared",
            uuid::Uuid::new_v4()
        );
        let db = Database::new(&url, "", 5);
        db.run_migrations()
            .await
            .expect("migrations should succeed");
        DAL::new(db)
    }

    /// Helper: create a workflow execution and return its ID.
    #[cfg(feature = "sqlite")]
    async fn create_workflow(dal: &DAL) -> UniversalUuid {
        dal.workflow_execution()
            .create(NewWorkflowExecution {
                workflow_name: "recovery_workflow".into(),
                workflow_version: "1.0".into(),
                status: "Running".into(),
                context_id: None,
            })
            .await
            .unwrap()
            .id
    }

    /// Helper: create a task with a given status, returning its ID.
    #[cfg(feature = "sqlite")]
    async fn create_task(
        dal: &DAL,
        workflow_id: UniversalUuid,
        name: &str,
        status: &str,
        attempt: i32,
        max_attempts: i32,
    ) -> UniversalUuid {
        dal.task_execution()
            .create(NewTaskExecution {
                workflow_execution_id: workflow_id,
                task_name: name.into(),
                status: status.into(),
                attempt,
                max_attempts,
                trigger_rules: r#"{"type":"Always"}"#.into(),
                task_configuration: "{}".into(),
            })
            .await
            .unwrap()
            .id
    }

    // ── reset_task_for_recovery ────────────────────────────────────

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_reset_task_for_recovery() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        let task_id = create_task(&dal, workflow_id, "recover_me", "Running", 1, 3).await;

        dal.task_execution()
            .reset_task_for_recovery(task_id)
            .await
            .unwrap();

        let task = dal.task_execution().get_by_id(task_id).await.unwrap();
        assert_eq!(task.status, "Ready");
        assert!(task.started_at.is_none());
        assert_eq!(task.recovery_attempts, 1);
        assert!(task.last_recovery_at.is_some());
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_reset_task_increments_recovery_attempts() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        let task_id = create_task(&dal, workflow_id, "multi_recover", "Running", 1, 3).await;

        // First recovery
        dal.task_execution()
            .reset_task_for_recovery(task_id)
            .await
            .unwrap();
        let task = dal.task_execution().get_by_id(task_id).await.unwrap();
        assert_eq!(task.recovery_attempts, 1);

        // Simulate it going back to Running (would normally happen via claiming)
        // We can just reset again since the method only checks the ID
        dal.task_execution()
            .reset_task_for_recovery(task_id)
            .await
            .unwrap();
        let task = dal.task_execution().get_by_id(task_id).await.unwrap();
        assert_eq!(task.recovery_attempts, 2);
    }

    // ── check_workflow_failure ─────────────────────────────────────

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_check_workflow_failure_no_abandoned() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        create_task(&dal, workflow_id, "ok", "Completed", 1, 3).await;

        let failed = dal
            .task_execution()
            .check_workflow_failure(workflow_id)
            .await
            .unwrap();
        assert!(!failed);
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_check_workflow_failure_with_abandoned() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        let task_id = create_task(&dal, workflow_id, "abandoned", "NotStarted", 1, 3).await;

        // Mark the task as abandoned (which sets status=Failed + error_details=ABANDONED:...)
        dal.task_execution()
            .mark_abandoned(task_id, "worker lost")
            .await
            .unwrap();

        let failed = dal
            .task_execution()
            .check_workflow_failure(workflow_id)
            .await
            .unwrap();
        assert!(failed);
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_check_workflow_failure_regular_failure_not_abandoned() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        let task_id = create_task(&dal, workflow_id, "regular_fail", "NotStarted", 1, 3).await;

        // A regular failure (not ABANDONED) should NOT trigger workflow failure check
        dal.task_execution()
            .mark_failed(task_id, "something broke", None)
            .await
            .unwrap();

        let failed = dal
            .task_execution()
            .check_workflow_failure(workflow_id)
            .await
            .unwrap();
        assert!(!failed);
    }

    // ── get_retry_stats ────────────────────────────────────────────

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_get_retry_stats_no_retries() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        create_task(&dal, workflow_id, "t1", "Completed", 1, 3).await;
        create_task(&dal, workflow_id, "t2", "Completed", 1, 3).await;

        let stats = dal
            .task_execution()
            .get_retry_stats(workflow_id)
            .await
            .unwrap();
        assert_eq!(stats.tasks_with_retries, 0);
        assert_eq!(stats.total_retries, 0);
        assert_eq!(stats.max_attempts_used, 1);
        assert_eq!(stats.tasks_exhausted_retries, 0);
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_get_retry_stats_with_retries() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;

        // Task that succeeded on attempt 1
        create_task(&dal, workflow_id, "first_try", "Completed", 1, 3).await;
        // Task that succeeded on attempt 3
        create_task(&dal, workflow_id, "third_try", "Completed", 3, 3).await;
        // Task that exhausted retries
        create_task(&dal, workflow_id, "exhausted", "Failed", 3, 3).await;

        let stats = dal
            .task_execution()
            .get_retry_stats(workflow_id)
            .await
            .unwrap();
        assert_eq!(stats.tasks_with_retries, 2); // third_try + exhausted
        assert_eq!(stats.total_retries, 4); // (3-1) + (3-1) = 4
        assert_eq!(stats.max_attempts_used, 3);
        assert_eq!(stats.tasks_exhausted_retries, 1); // exhausted
    }

    // ── get_exhausted_retry_tasks ──────────────────────────────────

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_get_exhausted_retry_tasks() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;

        create_task(&dal, workflow_id, "ok", "Completed", 1, 3).await;
        create_task(&dal, workflow_id, "still_trying", "Failed", 2, 3).await;
        let exhausted_id = create_task(&dal, workflow_id, "gave_up", "Failed", 3, 3).await;

        let exhausted = dal
            .task_execution()
            .get_exhausted_retry_tasks(workflow_id)
            .await
            .unwrap();
        assert_eq!(exhausted.len(), 1);
        assert_eq!(exhausted[0].id, exhausted_id);
        assert_eq!(exhausted[0].task_name, "gave_up");
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn test_get_exhausted_retry_tasks_empty() {
        let dal = unique_dal().await;
        let workflow_id = create_workflow(&dal).await;
        create_task(&dal, workflow_id, "ok", "Completed", 1, 3).await;

        let exhausted = dal
            .task_execution()
            .get_exhausted_retry_tasks(workflow_id)
            .await
            .unwrap();
        assert!(exhausted.is_empty());
    }
}