use std::collections::HashMap;

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

use crate::config::StageConfig;
use crate::error::{Result, WorkflowError};

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StageState {
    pub name: String,
    pub status: StageStatus,
    pub config: StageConfig,
    pub assigned_tasks: HashMap<String, TaskState>,
    pub artifacts_produced: Vec<String>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskState {
    pub task_id: String,
    pub agent_id: String,
    pub agent_role: String,
    pub status: TaskStatus,
    pub result: Option<String>,
    pub artifacts: Vec<String>,
    pub assigned_at: DateTime<Utc>,
    pub completed_at: Option<DateTime<Utc>>,
}

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

impl StageState {
    pub fn new(config: StageConfig) -> Self {
        Self {
            name: config.name.clone(),
            status: StageStatus::Pending,
            config,
            assigned_tasks: HashMap::new(),
            artifacts_produced: Vec::new(),
            started_at: None,
            completed_at: None,
        }
    }

    pub fn all_tasks_completed(&self) -> bool {
        !self.assigned_tasks.is_empty()
            && self
                .assigned_tasks
                .values()
                .all(|t| matches!(t.status, TaskStatus::Completed))
    }

    pub fn any_task_failed(&self) -> bool {
        self.assigned_tasks
            .values()
            .any(|t| matches!(t.status, TaskStatus::Failed(_)))
    }

    pub fn collect_artifacts(&self) -> Vec<String> {
        self.assigned_tasks
            .values()
            .flat_map(|t| t.artifacts.clone())
            .collect()
    }

    pub fn transition(&mut self, new_status: StageStatus) -> Result<()> {
        let valid = matches!(
            (&self.status, &new_status),
            (StageStatus::Pending, StageStatus::Running)
                | (StageStatus::Pending, StageStatus::WaitingApproval)
                | (StageStatus::WaitingApproval, StageStatus::Running)
                | (StageStatus::Running, StageStatus::Completed)
                | (StageStatus::Running, StageStatus::Failed(_))
        );

        if !valid {
            return Err(WorkflowError::InvalidTransition {
                from: format!("{:?}", self.status),
                to: format!("{:?}", new_status),
            });
        }

        self.status = new_status;
        Ok(())
    }

    pub fn mark_started(&mut self) {
        self.started_at = Some(Utc::now());
    }

    pub fn mark_completed(&mut self) {
        self.completed_at = Some(Utc::now());
    }
}

impl TaskState {
    pub fn new(task_id: String, agent_id: String, agent_role: String) -> Self {
        Self {
            task_id,
            agent_id,
            agent_role,
            status: TaskStatus::Assigned,
            result: None,
            artifacts: Vec::new(),
            assigned_at: Utc::now(),
            completed_at: None,
        }
    }

    pub fn mark_running(&mut self) {
        self.status = TaskStatus::Running;
    }

    pub fn mark_completed(&mut self, result: String, artifacts: Vec<String>) {
        self.status = TaskStatus::Completed;
        self.result = Some(result);
        self.artifacts = artifacts;
        self.completed_at = Some(Utc::now());
    }

    pub fn mark_failed(&mut self, error: String) {
        self.status = TaskStatus::Failed(error);
        self.completed_at = Some(Utc::now());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_config() -> StageConfig {
        StageConfig {
            name: "test_stage".to_string(),
            agents: vec!["agent1".to_string()],
            parallel: false,
            max_parallel: None,
            approval_required: false,
            compensation_agents: None,
        }
    }

    #[test]
    fn test_stage_state_creation() {
        let config = create_test_config();
        let stage = StageState::new(config.clone());

        assert_eq!(stage.name, "test_stage");
        assert_eq!(stage.status, StageStatus::Pending);
        assert!(stage.assigned_tasks.is_empty());
    }

    #[test]
    fn test_stage_transitions() {
        let config = create_test_config();
        let mut stage = StageState::new(config);

        assert!(stage.transition(StageStatus::Running).is_ok());
        assert_eq!(stage.status, StageStatus::Running);

        assert!(stage.transition(StageStatus::Completed).is_ok());
        assert_eq!(stage.status, StageStatus::Completed);
    }

    #[test]
    fn test_invalid_stage_transition() {
        let config = create_test_config();
        let mut stage = StageState::new(config);

        let result = stage.transition(StageStatus::Completed);
        assert!(result.is_err());
    }

    #[test]
    fn test_all_tasks_completed() {
        let config = create_test_config();
        let mut stage = StageState::new(config);

        let mut task1 = TaskState::new(
            "task1".to_string(),
            "agent1".to_string(),
            "developer".to_string(),
        );
        task1.mark_completed("done".to_string(), vec![]);

        let mut task2 = TaskState::new(
            "task2".to_string(),
            "agent2".to_string(),
            "developer".to_string(),
        );
        task2.mark_completed("done".to_string(), vec![]);

        stage.assigned_tasks.insert("task1".to_string(), task1);
        stage.assigned_tasks.insert("task2".to_string(), task2);

        assert!(stage.all_tasks_completed());
    }

    #[test]
    fn test_any_task_failed() {
        let config = create_test_config();
        let mut stage = StageState::new(config);

        let mut task1 = TaskState::new(
            "task1".to_string(),
            "agent1".to_string(),
            "developer".to_string(),
        );
        task1.mark_completed("done".to_string(), vec![]);

        let mut task2 = TaskState::new(
            "task2".to_string(),
            "agent2".to_string(),
            "developer".to_string(),
        );
        task2.mark_failed("error occurred".to_string());

        stage.assigned_tasks.insert("task1".to_string(), task1);
        stage.assigned_tasks.insert("task2".to_string(), task2);

        assert!(stage.any_task_failed());
    }

    #[test]
    fn test_collect_artifacts() {
        let config = create_test_config();
        let mut stage = StageState::new(config);

        let mut task1 = TaskState::new(
            "task1".to_string(),
            "agent1".to_string(),
            "developer".to_string(),
        );
        task1.mark_completed("done".to_string(), vec!["artifact1".to_string()]);

        let mut task2 = TaskState::new(
            "task2".to_string(),
            "agent2".to_string(),
            "developer".to_string(),
        );
        task2.mark_completed(
            "done".to_string(),
            vec!["artifact2".to_string(), "artifact3".to_string()],
        );

        stage.assigned_tasks.insert("task1".to_string(), task1);
        stage.assigned_tasks.insert("task2".to_string(), task2);

        let artifacts = stage.collect_artifacts();
        assert_eq!(artifacts.len(), 3);
        assert!(artifacts.contains(&"artifact1".to_string()));
        assert!(artifacts.contains(&"artifact2".to_string()));
        assert!(artifacts.contains(&"artifact3".to_string()));
    }

    #[test]
    fn test_task_state_lifecycle() {
        let mut task = TaskState::new(
            "task1".to_string(),
            "agent1".to_string(),
            "developer".to_string(),
        );

        assert_eq!(task.status, TaskStatus::Assigned);

        task.mark_running();
        assert_eq!(task.status, TaskStatus::Running);

        task.mark_completed("success".to_string(), vec!["artifact1".to_string()]);
        assert_eq!(task.status, TaskStatus::Completed);
        assert_eq!(task.result, Some("success".to_string()));
        assert_eq!(task.artifacts.len(), 1);
        assert!(task.completed_at.is_some());
    }

    #[test]
    fn test_task_failure() {
        let mut task = TaskState::new(
            "task1".to_string(),
            "agent1".to_string(),
            "developer".to_string(),
        );

        task.mark_running();
        task.mark_failed("compilation error".to_string());

        assert!(matches!(task.status, TaskStatus::Failed(_)));
        assert!(task.completed_at.is_some());
    }
}
