use std::sync::Arc;

use tracing::{error, info};
use vapora_swarm::coordinator::SwarmCoordinator;

use crate::error::Result;
use crate::stage::StageState;

/// Saga compensator that dispatches rollback tasks for completed stages when
/// a downstream stage fails.
///
/// Iterates executed stages in reverse order, and for each stage that has
/// `compensation_agents` configured, submits a `compensation` task to the
/// `SwarmCoordinator`. Errors are logged but never propagated — compensation
/// is best-effort (consistent with stratum-orchestrator's behaviour).
pub struct SagaCompensator {
    swarm: Arc<SwarmCoordinator>,
}

impl SagaCompensator {
    pub fn new(swarm: Arc<SwarmCoordinator>) -> Self {
        Self { swarm }
    }

    /// Dispatch compensation tasks for `executed_stages` in reverse order.
    ///
    /// `context` is the workflow's `initial_context` — passed to each
    /// compensation task so the agent has enough information to undo the work.
    pub async fn compensate(
        &self,
        workflow_id: &str,
        executed_stages: &[StageState],
        context: &serde_json::Value,
    ) -> Result<()> {
        for stage in executed_stages.iter().rev() {
            let compensation_agents = match &stage.config.compensation_agents {
                Some(agents) if !agents.is_empty() => agents.clone(),
                _ => continue,
            };

            let task_id = format!("compensate-{}-{}", workflow_id, stage.name);
            let task_payload = serde_json::json!({
                "type": "compensation",
                "stage_name": stage.name,
                "workflow_id": workflow_id,
                "original_context": context,
                "artifacts_to_undo": stage.artifacts_produced,
            });

            match self
                .swarm
                .submit_task_for_bidding(
                    task_id.clone(),
                    task_payload.to_string(),
                    compensation_agents,
                )
                .await
            {
                Ok(_) => {
                    info!(
                        workflow_id = %workflow_id,
                        stage = %stage.name,
                        task_id = %task_id,
                        "Saga compensation task dispatched"
                    );
                }
                Err(e) => {
                    error!(
                        workflow_id = %workflow_id,
                        stage = %stage.name,
                        task_id = %task_id,
                        error = %e,
                        "Compensation dispatch failed (best-effort, continuing)"
                    );
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::config::StageConfig;
    use crate::stage::{StageState, StageStatus, TaskState};

    fn stage_with_compensation(name: &str, agents: Vec<String>) -> StageState {
        let config = StageConfig {
            name: name.to_string(),
            agents: vec!["developer".to_string()],
            parallel: false,
            max_parallel: None,
            approval_required: false,
            compensation_agents: Some(agents),
        };
        let mut state = StageState::new(config);
        state.status = StageStatus::Completed;
        state.artifacts_produced = vec![format!("{name}-artifact")];
        state
    }

    fn stage_without_compensation(name: &str) -> StageState {
        let config = StageConfig {
            name: name.to_string(),
            agents: vec!["developer".to_string()],
            parallel: false,
            max_parallel: None,
            approval_required: false,
            compensation_agents: None,
        };
        let mut state = StageState::new(config);
        state.status = StageStatus::Completed;
        state
    }

    #[test]
    fn test_stages_with_no_compensation_agents_are_skipped() {
        let stages = vec![
            stage_without_compensation("design"),
            stage_without_compensation("review"),
        ];

        // Verify the filter logic: stages without compensation_agents produce no tasks
        let compensatable: Vec<&StageState> = stages
            .iter()
            .rev()
            .filter(|s| {
                s.config
                    .compensation_agents
                    .as_ref()
                    .map(|a| !a.is_empty())
                    .unwrap_or(false)
            })
            .collect();

        assert!(compensatable.is_empty());
    }

    #[test]
    fn test_stages_with_compensation_agents_are_included() {
        let stages = vec![
            stage_with_compensation("deploy", vec!["ops-agent".to_string()]),
            stage_without_compensation("review"),
            stage_with_compensation("provision", vec!["infra-agent".to_string()]),
        ];

        let compensatable: Vec<&str> = stages
            .iter()
            .rev()
            .filter(|s| {
                s.config
                    .compensation_agents
                    .as_ref()
                    .map(|a| !a.is_empty())
                    .unwrap_or(false)
            })
            .map(|s| s.name.as_str())
            .collect();

        // Reverse order: provision first, then deploy
        assert_eq!(compensatable, vec!["provision", "deploy"]);
    }

    #[test]
    fn test_task_state_mark_completed_carries_artifacts() {
        let mut task = TaskState::new(
            "t1".to_string(),
            "agent-1".to_string(),
            "ops-agent".to_string(),
        );
        task.mark_completed("done".to_string(), vec!["artifact-1".to_string()]);
        assert_eq!(task.artifacts, vec!["artifact-1"]);
    }
}
