use std::collections::HashMap;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::artifact::Artifact;
use crate::config::WorkflowConfig;
use crate::error::Result;
use crate::stage::{StageState, StageStatus};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkflowStatus {
    Running,
    WaitingApproval(usize),
    Completed,
    Failed(String),
    Cancelled,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowInstance {
    pub id: String,
    pub template_name: String,
    pub status: WorkflowStatus,
    pub stages: Vec<StageState>,
    pub current_stage_idx: usize,
    pub initial_context: Value,
    pub accumulated_artifacts: HashMap<String, Artifact>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

impl WorkflowInstance {
    pub fn new(template: &WorkflowConfig, initial_context: Value) -> Self {
        let stages = template
            .stages
            .iter()
            .map(|sc| StageState::new(sc.clone()))
            .collect();

        Self {
            id: uuid::Uuid::new_v4().to_string(),
            template_name: template.name.clone(),
            status: WorkflowStatus::Running,
            stages,
            current_stage_idx: 0,
            initial_context,
            accumulated_artifacts: HashMap::new(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
        }
    }

    pub fn current_stage(&self) -> Option<&StageState> {
        self.stages.get(self.current_stage_idx)
    }

    pub fn current_stage_mut(&mut self) -> Option<&mut StageState> {
        self.stages.get_mut(self.current_stage_idx)
    }

    pub fn advance_stage(&mut self) -> Result<bool> {
        if let Some(stage) = self.current_stage() {
            for artifact_id in stage.collect_artifacts() {
                if let Some(artifact) = self.accumulated_artifacts.get(&artifact_id) {
                    tracing::debug!(
                        artifact_id = %artifact_id,
                        artifact_type = ?artifact.artifact_type,
                        "Artifact collected from stage"
                    );
                }
            }
        }

        self.current_stage_idx += 1;
        self.updated_at = Utc::now();

        if self.current_stage_idx >= self.stages.len() {
            self.status = WorkflowStatus::Completed;
            Ok(false)
        } else {
            Ok(true)
        }
    }

    pub fn build_stage_context(&self) -> Value {
        let mut ctx = self.initial_context.clone();

        if let Some(obj) = ctx.as_object_mut() {
            obj.insert(
                "artifacts".to_string(),
                serde_json::to_value(&self.accumulated_artifacts).unwrap_or_default(),
            );
            obj.insert("workflow_id".to_string(), Value::String(self.id.clone()));
            obj.insert(
                "stage_idx".to_string(),
                Value::Number(self.current_stage_idx.into()),
            );

            if let Some(stage) = self.current_stage() {
                obj.insert("stage_name".to_string(), Value::String(stage.name.clone()));
            }

            obj.insert(
                "previous_stages".to_string(),
                serde_json::to_value(
                    self.stages[..self.current_stage_idx]
                        .iter()
                        .map(|s| {
                            serde_json::json!({
                                "name": s.name,
                                "status": s.status,
                                "artifacts": s.artifacts_produced,
                            })
                        })
                        .collect::<Vec<_>>(),
                )
                .unwrap_or_default(),
            );
        }

        ctx
    }

    pub fn add_artifact(&mut self, artifact: Artifact) {
        let artifact_id = artifact.id.clone();
        self.accumulated_artifacts
            .insert(artifact_id.clone(), artifact);

        if let Some(stage) = self.current_stage_mut() {
            if !stage.artifacts_produced.contains(&artifact_id) {
                stage.artifacts_produced.push(artifact_id);
            }
        }

        self.updated_at = Utc::now();
    }

    pub fn is_completed(&self) -> bool {
        matches!(self.status, WorkflowStatus::Completed)
    }

    pub fn is_failed(&self) -> bool {
        matches!(self.status, WorkflowStatus::Failed(_))
    }

    pub fn is_cancelled(&self) -> bool {
        matches!(self.status, WorkflowStatus::Cancelled)
    }

    pub fn is_waiting_approval(&self) -> bool {
        matches!(self.status, WorkflowStatus::WaitingApproval(_))
    }

    pub fn cancel(&mut self, reason: String) {
        self.status = WorkflowStatus::Cancelled;
        self.updated_at = Utc::now();

        if let Some(stage) = self.current_stage_mut() {
            let _ = stage.transition(StageStatus::Failed(reason));
        }
    }

    pub fn fail(&mut self, error: String) {
        self.status = WorkflowStatus::Failed(error.clone());
        self.updated_at = Utc::now();

        if let Some(stage) = self.current_stage_mut() {
            let _ = stage.transition(StageStatus::Failed(error));
        }
    }

    /// Marks a task as failed in the current stage and optionally transitions
    /// the stage to `Failed`.
    ///
    /// Returns `Some(Ok(stage_name))` when the failure is non-retryable and
    /// the stage was transitioned, `None` when retrying or no active stage.
    /// Isolates the mutable borrow of `current_stage_mut()` so callers can
    /// access other fields of `self` after this returns.
    pub fn mark_current_task_failed(
        &mut self,
        task_id: &str,
        error: &str,
        can_retry: bool,
    ) -> Option<Result<String>> {
        let stage = self.current_stage_mut()?;
        if let Some(task) = stage.assigned_tasks.get_mut(task_id) {
            task.mark_failed(error.to_string());
        }
        if can_retry {
            return None;
        }
        let stage_name = stage.name.clone();
        let r = stage.transition(StageStatus::Failed(error.to_string()));
        Some(r.map(|()| stage_name))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::artifact::ArtifactType;
    use crate::config::StageConfig;

    fn create_test_workflow_config() -> WorkflowConfig {
        WorkflowConfig {
            name: "test_workflow".to_string(),
            trigger: "manual".to_string(),
            stages: vec![
                StageConfig {
                    name: "stage1".to_string(),
                    agents: vec!["agent1".to_string()],
                    parallel: false,
                    max_parallel: None,
                    approval_required: false,
                    compensation_agents: None,
                },
                StageConfig {
                    name: "stage2".to_string(),
                    agents: vec!["agent2".to_string()],
                    parallel: false,
                    max_parallel: None,
                    approval_required: false,
                    compensation_agents: None,
                },
            ],
            schedule: None,
            notifications: Default::default(),
        }
    }

    #[test]
    fn test_workflow_instance_creation() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({"key": "value"});
        let instance = WorkflowInstance::new(&config, context);

        assert_eq!(instance.template_name, "test_workflow");
        assert_eq!(instance.status, WorkflowStatus::Running);
        assert_eq!(instance.stages.len(), 2);
        assert_eq!(instance.current_stage_idx, 0);
    }

    #[test]
    fn test_current_stage() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let instance = WorkflowInstance::new(&config, context);

        let current = instance.current_stage().unwrap();
        assert_eq!(current.name, "stage1");
    }

    #[test]
    fn test_advance_stage() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let mut instance = WorkflowInstance::new(&config, context);

        let has_more = instance.advance_stage().unwrap();
        assert!(has_more);
        assert_eq!(instance.current_stage_idx, 1);
        assert_eq!(instance.current_stage().unwrap().name, "stage2");

        let has_more = instance.advance_stage().unwrap();
        assert!(!has_more);
        assert_eq!(instance.status, WorkflowStatus::Completed);
    }

    #[test]
    fn test_build_stage_context() {
        let config = create_test_workflow_config();
        let initial_context = serde_json::json!({"initial": "data"});
        let instance = WorkflowInstance::new(&config, initial_context);

        let ctx = instance.build_stage_context();
        assert_eq!(ctx["initial"], "data");
        assert_eq!(ctx["workflow_id"], instance.id);
        assert_eq!(ctx["stage_idx"], 0);
        assert_eq!(ctx["stage_name"], "stage1");
    }

    #[test]
    fn test_add_artifact() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let mut instance = WorkflowInstance::new(&config, context);

        let artifact = Artifact::new(
            ArtifactType::Code,
            "stage1".to_string(),
            "agent1".to_string(),
            "content_ref".to_string(),
        );

        let artifact_id = artifact.id.clone();
        instance.add_artifact(artifact);

        assert!(instance.accumulated_artifacts.contains_key(&artifact_id));
        assert!(instance
            .current_stage()
            .unwrap()
            .artifacts_produced
            .contains(&artifact_id));
    }

    #[test]
    fn test_workflow_status_checks() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let mut instance = WorkflowInstance::new(&config, context);

        assert!(!instance.is_completed());
        assert!(!instance.is_failed());
        assert!(!instance.is_cancelled());

        instance.status = WorkflowStatus::Completed;
        assert!(instance.is_completed());

        instance.status = WorkflowStatus::Failed("error".to_string());
        assert!(instance.is_failed());

        instance.status = WorkflowStatus::Cancelled;
        assert!(instance.is_cancelled());
    }

    #[test]
    fn test_cancel_workflow() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let mut instance = WorkflowInstance::new(&config, context);

        instance.cancel("User requested".to_string());

        assert!(instance.is_cancelled());
    }

    #[test]
    fn test_fail_workflow() {
        let config = create_test_workflow_config();
        let context = serde_json::json!({});
        let mut instance = WorkflowInstance::new(&config, context);

        instance.fail("Task execution failed".to_string());

        assert!(instance.is_failed());
    }
}
