use std::collections::HashMap;
use uuid::Uuid;
use common_contracts::workflow_types::{TaskStatus, TaskContext};
use common_contracts::messages::{TaskType, TaskMetadata};
use workflow_context::{Vgcs, ContextStore};
use anyhow::Result;
use tracing::info;
use serde::{Serialize, Deserialize};

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TaskExecutionBuffer {
    pub logs: Vec<String>,
    pub content_buffer: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommitTracker {
    /// Maps task_id to the commit hash it produced.
    pub task_commits: HashMap<String, String>,
    /// The latest merged commit for the whole workflow (if linear).
    /// Or just a reference to the "main" branch tip.
    pub head_commit: String,
    /// Maps task_id to generic metadata (e.g. output_path) (New)
    pub task_metadata: HashMap<String, TaskMetadata>,
}

impl CommitTracker {
    pub fn new(initial_commit: String) -> Self {
        Self {
            task_commits: HashMap::new(),
            head_commit: initial_commit,
            task_metadata: HashMap::new(),
        }
    }

    pub fn record_commit(&mut self, task_id: &str, commit: String) {
        self.task_commits.insert(task_id.to_string(), commit.clone());
        // Note: head_commit update strategy depends on whether we want to track 
        // a single "main" branch or just use task_commits for DAG resolution.
        // For now, we don't eagerly update head_commit unless it's a final task.
    }
    
    pub fn record_metadata(&mut self, task_id: &str, meta: serde_json::Value) {
        // Convert generic JSON to TaskMetadata
        // If the incoming meta already has execution_log_path/output_path at top level, 
        // we should extract them to avoid duplication in 'extra' if they are also left there.
        
        let mut execution_log_path = None;
        let mut output_path = None;
        let mut extra = HashMap::new();

        if let Some(obj) = meta.as_object() {
            for (k, v) in obj {
                if k == "execution_log_path" {
                    execution_log_path = v.as_str().map(|s| s.to_string());
                } else if k == "output_path" {
                    output_path = v.as_str().map(|s| s.to_string());
                } else {
                    // Avoid nested "extra" if the input was already structured
                    if k == "extra" && v.is_object() {
                         if let Some(inner_extra) = v.as_object() {
                             for (ik, iv) in inner_extra {
                                 extra.insert(ik.clone(), iv.clone());
                             }
                         }
                    } else {
                        extra.insert(k.clone(), v.clone());
                    }
                }
            }
        }

        let metadata = TaskMetadata {
            output_path,
            execution_log_path,
            extra,
        };
        self.task_metadata.insert(task_id.to_string(), metadata);
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DagScheduler {
    pub request_id: Uuid,
    pub nodes: HashMap<String, DagNode>,
    /// TaskID -> List of downstream TaskIDs
    pub forward_deps: HashMap<String, Vec<String>>,
    /// TaskID -> List of upstream TaskIDs
    pub reverse_deps: HashMap<String, Vec<String>>,
    
    pub commit_tracker: CommitTracker,
    
    /// Flag to track if the workflow completion event has already been emitted.
    #[serde(default)]
    pub workflow_finished_flag: bool,
    
    #[serde(default = "default_start_time")]
    pub start_time: chrono::DateTime<chrono::Utc>,

    /// In-memory buffer for active tasks' logs and content
    #[serde(default)]
    pub task_execution_states: HashMap<String, TaskExecutionBuffer>,
}

fn default_start_time() -> chrono::DateTime<chrono::Utc> {
    chrono::Utc::now()
}

impl DagScheduler {
    pub fn to_dto(&self) -> common_contracts::messages::WorkflowDag {
        let edges = self.forward_deps.iter().flat_map(|(from, tos)| {
            tos.iter().map(move |to| common_contracts::messages::TaskDependency {
                from: from.clone(),
                to: to.clone(),
            })
        }).collect();

        let nodes = self.nodes.values().map(|n| common_contracts::messages::TaskNode {
            id: n.id.clone(),
            name: n.id.clone(), // Use ID as name for now, or add name field to DagNode
            display_name: n.display_name.clone(),
            r#type: n.task_type,
            initial_status: match n.status {
                TaskStatus::Pending => common_contracts::messages::TaskStatus::Pending,
                TaskStatus::Scheduled => common_contracts::messages::TaskStatus::Scheduled,
                TaskStatus::Running => common_contracts::messages::TaskStatus::Running,
                TaskStatus::Completed => common_contracts::messages::TaskStatus::Completed,
                TaskStatus::Failed => common_contracts::messages::TaskStatus::Failed,
                TaskStatus::Skipped => common_contracts::messages::TaskStatus::Skipped,
                TaskStatus::Cancelled => common_contracts::messages::TaskStatus::Skipped,
            },
        }).collect();

        common_contracts::messages::WorkflowDag {
            nodes,
            edges,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DagNode {
    pub id: String,
    pub display_name: Option<String>,
    pub task_type: TaskType, // Kept for UI/Observability, not for logic
    pub status: TaskStatus,
    pub config: serde_json::Value,
    pub routing_key: String,
    /// The commit hash used as input for this task
    pub input_commit: Option<String>,
    
    // --- Observability & Watchdog ---
    pub started_at: Option<i64>, // Timestamp ms
    pub last_heartbeat_at: Option<i64>, // Timestamp ms
}


impl DagScheduler {
    pub fn new(request_id: Uuid, initial_commit: String) -> Self {
        Self {
            request_id,
            nodes: HashMap::new(),
            forward_deps: HashMap::new(),
            reverse_deps: HashMap::new(),
            commit_tracker: CommitTracker::new(initial_commit),
            workflow_finished_flag: false,
            start_time: chrono::Utc::now(),
            task_execution_states: HashMap::new(),
        }
    }

    pub fn add_node(&mut self, id: String, display_name: Option<String>, task_type: TaskType, routing_key: String, config: serde_json::Value) {
        self.nodes.insert(id.clone(), DagNode {
            id,
            display_name,
            task_type,
            status: TaskStatus::Pending,
            config,
            routing_key,
            input_commit: None,
            started_at: None,
            last_heartbeat_at: None,
        });
    }

    pub fn set_input_commit(&mut self, task_id: &str, commit: String) {
        if let Some(node) = self.nodes.get_mut(task_id) {
            node.input_commit = Some(commit);
        }
    }

    pub fn add_dependency(&mut self, from: &str, to: &str) {
        self.forward_deps.entry(from.to_string()).or_default().push(to.to_string());
        self.reverse_deps.entry(to.to_string()).or_default().push(from.to_string());
    }

    /// Get all tasks that have no dependencies (roots)
    pub fn get_initial_tasks(&self) -> Vec<String> {
        self.nodes.values()
            .filter(|n| n.status == TaskStatus::Pending && self.reverse_deps.get(&n.id).map_or(true, |deps| deps.is_empty()))
            .map(|n| n.id.clone())
            .collect()
    }

    pub fn update_status(&mut self, task_id: &str, status: TaskStatus) {
        if let Some(node) = self.nodes.get_mut(task_id) {
            // State transition logic for timestamps
            if node.status != TaskStatus::Running && status == TaskStatus::Running {
                 // Transition to Running: Set start time if not set (or reset?)
                 // Actually dispatch sets to Scheduled. Worker sets to Running.
                 // Let's assume Running means "Worker Started".
                 if node.started_at.is_none() {
                     node.started_at = Some(chrono::Utc::now().timestamp_millis());
                 }
                 node.last_heartbeat_at = Some(chrono::Utc::now().timestamp_millis());
            }
            
            // If we are scheduling, we might want to set a "dispatch_time" but let's use started_at for simplicity
            // or maybe we strictly follow: Running = Started.
            
            node.status = status;
        }
    }
    
    pub fn update_heartbeat(&mut self, task_id: &str, timestamp: i64) {
        if let Some(node) = self.nodes.get_mut(task_id) {
            node.last_heartbeat_at = Some(timestamp);
        }
    }
    
    pub fn get_status(&self, task_id: &str) -> TaskStatus {
        self.nodes.get(task_id).map(|n| n.status).unwrap_or(TaskStatus::Pending)
    }
    
    pub fn record_result(&mut self, task_id: &str, new_commit: Option<String>) {
        if let Some(c) = new_commit {
            self.commit_tracker.record_commit(task_id, c);
        }
    }

    pub fn record_metadata(&mut self, task_id: &str, meta: serde_json::Value) {
        self.commit_tracker.record_metadata(task_id, meta);
    }

    pub fn append_log(&mut self, task_id: &str, log: String) {
        self.task_execution_states
            .entry(task_id.to_string())
            .or_default()
            .logs
            .push(log);
    }

    pub fn append_content(&mut self, task_id: &str, content_delta: &str) {
        self.task_execution_states
            .entry(task_id.to_string())
            .or_default()
            .content_buffer
            .push_str(content_delta);
    }

    /// Recursively cancel downstream tasks of a failed/skipped task.
    pub fn cancel_downstream(&mut self, task_id: &str) {
        // Clone deps to avoid borrowing self while mutating nodes
        if let Some(downstream) = self.forward_deps.get(task_id).cloned() {
            for next_id in downstream {
                let should_recurse = if let Some(node) = self.nodes.get_mut(&next_id) {
                    if node.status == TaskStatus::Pending || node.status == TaskStatus::Scheduled {
                        node.status = TaskStatus::Skipped; // Use Skipped for cascading cancellation
                        true
                    } else {
                        false
                    }
                } else {
                    false
                };

                if should_recurse {
                    self.cancel_downstream(&next_id);
                }
            }
        }
    }

    /// Check if all tasks in the DAG have reached a terminal state.
    pub fn is_workflow_finished(&self) -> bool {
        self.nodes.values().all(|n| matches!(n.status, 
            TaskStatus::Completed | 
            TaskStatus::Failed | 
            TaskStatus::Skipped | 
            TaskStatus::Cancelled
        ))
    }

    /// Checks if the workflow is finished AND not yet marked as finished.
    /// Returns true only if this is the FIRST time the workflow is detected as finished.
    /// Updates the internal flag.
    pub fn try_finish_workflow(&mut self) -> bool {
        if self.workflow_finished_flag {
            return false;
        }
        
        if self.is_workflow_finished() {
            self.workflow_finished_flag = true;
            return true;
        }
        
        false
    }

    /// Check if any task has failed, indicating the workflow is partially or fully failed.
    /// Note: Depending on requirements, some failures might be tolerant. 
    /// Here we assume any failure means the workflow has failed components.
    pub fn has_failures(&self) -> bool {
        self.nodes.values().any(|n| n.status == TaskStatus::Failed)
    }

    /// Determine which tasks are ready to run given that `completed_task_id` just finished.
    pub fn get_ready_downstream_tasks(&self, completed_task_id: &str) -> Vec<String> {
        let mut ready: Vec<String> = Vec::new();
        if let Some(downstream) = self.forward_deps.get(completed_task_id) {
            for next_id in downstream {
                if self.is_ready(next_id) {
                    ready.push(next_id.clone());
                }
            }
        }
        ready
    }

    fn is_ready(&self, task_id: &str) -> bool {
        let node = match self.nodes.get(task_id) {
            Some(n) => n,
            None => return false,
        };

        if node.status != TaskStatus::Pending {
            return false;
        }

        if let Some(deps) = self.reverse_deps.get(task_id) {
            for dep_id in deps {
                match self.nodes.get(dep_id).map(|n| n.status) {
                    Some(TaskStatus::Completed) | Some(TaskStatus::Failed) | Some(TaskStatus::Skipped) | Some(TaskStatus::Cancelled) => continue,
                    _ => return false, // Dependency not finished
                }
            }
        }
        true
    }

    /// Resolve the context (Base Commit) for a task.
    /// If multiple dependencies, perform Merge or Fast-Forward.
    pub fn resolve_context(&self, task_id: &str, vgcs: &Vgcs) -> Result<TaskContext> {
        let deps = self.reverse_deps.get(task_id).cloned().unwrap_or_default();
        
        if deps.is_empty() {
            // Root task: Use initial commit (usually empty string or base snapshot)
            return Ok(TaskContext {
                base_commit: Some(self.commit_tracker.head_commit.clone()),
                mount_path: None,
            });
        }

        // Collect parent commits
        let mut parent_commits = Vec::new();
        for dep_id in &deps {
            if let Some(c) = self.commit_tracker.task_commits.get(dep_id) {
                if !c.is_empty() {
                    parent_commits.push(c.clone());
                }
            }
        }

        if parent_commits.is_empty() {
            // All parents produced no commit? Fallback to head or empty.
            return Ok(TaskContext {
                base_commit: Some(self.commit_tracker.head_commit.clone()),
                mount_path: None,
            });
        }

        // Merge Strategy
        let final_commit = self.merge_commits(vgcs, parent_commits)?;

        Ok(TaskContext {
            base_commit: Some(final_commit),
            mount_path: None, // Or determine based on config
        })
    }

    /// Merge logic:
    /// 1 parent -> Return it.
    /// 2+ parents -> Iteratively merge using smart merge_commits
    fn merge_commits(&self, vgcs: &Vgcs, commits: Vec<String>) -> Result<String> {
        if commits.is_empty() {
            return Ok(String::new());
        }
        if commits.len() == 1 {
            return Ok(commits[0].clone());
        }
        
        let mut current_head = commits[0].clone();
        
        for i in 1..commits.len() {
            let next_commit = &commits[i];
            if current_head == *next_commit {
                continue;
            }
            
            // Use the smart merge_commits which finds the common ancestor automatically
            // Note: This handles Fast-Forward implicitly (merge_commits checks for ancestry)
            info!("Merging commits: Ours={}, Theirs={}", current_head, next_commit);
            current_head = vgcs.merge_commits(&self.request_id.to_string(), &current_head, next_commit)?;
        }

        Ok(current_head)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;
    use workflow_context::{Vgcs, ContextStore};
    use common_contracts::messages::TaskType;
    use serde_json::json;

    #[test]
    fn test_dag_double_completion_prevention() {
        let req_id = Uuid::new_v4();
        let mut dag = DagScheduler::new(req_id, "init".to_string());
        
        dag.add_node("A".to_string(), None, TaskType::DataFetch, "key".into(), json!({}));
        
        // 1. Not finished
        assert!(!dag.is_workflow_finished());
        assert!(!dag.try_finish_workflow());
        
        // 2. Complete task
        dag.update_status("A", TaskStatus::Completed);
        assert!(dag.is_workflow_finished());
        
        // 3. Try finish - First time should be true
        assert!(dag.try_finish_workflow());
        assert!(dag.workflow_finished_flag);
        
        // 4. Try finish - Second time should be false (Idempotency check)
        assert!(!dag.try_finish_workflow());
    }

    #[test]
    fn test_dag_merge_strategy() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let vgcs = Vgcs::new(temp_dir.path());
        let req_id = Uuid::new_v4();
        let req_id_str = req_id.to_string();

        vgcs.init_repo(&req_id_str)?;

        // 0. Create Initial Commit (Common Ancestor)
        let tx = vgcs.begin_transaction(&req_id_str, "")?;
        let init_commit = Box::new(tx).commit("Initial Commit", "system")?;

        // 1. Setup DAG
        let mut dag = DagScheduler::new(req_id, init_commit.clone());
        dag.add_node("A".to_string(), None, TaskType::DataFetch, "key.a".into(), json!({}));
        dag.add_node("B".to_string(), None, TaskType::DataFetch, "key.b".into(), json!({}));
        dag.add_node("C".to_string(), None, TaskType::Analysis, "key.c".into(), json!({}));

        // C depends on A and B
        dag.add_dependency("A", "C");
        dag.add_dependency("B", "C");

        // 2. Simulate Task A Execution -> Commit A (Based on Init)
        let mut tx = vgcs.begin_transaction(&req_id_str, &init_commit)?;
        tx.write("file_a.txt", b"Content A")?;
        let commit_a = Box::new(tx).commit("Task A", "worker")?;
        dag.record_result("A", Some(commit_a.clone()));
        dag.update_status("A", TaskStatus::Completed);

        // 3. Simulate Task B Execution -> Commit B (Based on Init)
        let mut tx = vgcs.begin_transaction(&req_id_str, &init_commit)?;
        tx.write("file_b.txt", b"Content B")?;
        let commit_b = Box::new(tx).commit("Task B", "worker")?;
        dag.record_result("B", Some(commit_b.clone()));
        dag.update_status("B", TaskStatus::Completed);

        // 4. Resolve Context for C
        // Should merge A and B
        let ctx = dag.resolve_context("C", &vgcs)?;
        let merged_commit = ctx.base_commit.expect("Should have base commit");

        // Verify merged content
        let files = vgcs.list_dir(&req_id_str, &merged_commit, "")?;
        let file_names: Vec<String> = files.iter().map(|f| f.name.clone()).collect();
        
        assert!(file_names.contains(&"file_a.txt".to_string()));
        assert!(file_names.contains(&"file_b.txt".to_string()));
        
        Ok(())
    }

    #[test]
    fn test_dag_timestamp_updates() {
        let req_id = Uuid::new_v4();
        let mut dag = DagScheduler::new(req_id, "init".to_string());
        dag.add_node("A".to_string(), None, TaskType::DataFetch, "key".into(), json!({}));

        // 1. Initial state
        let node = dag.nodes.get("A").unwrap();
        assert_eq!(node.status, TaskStatus::Pending);
        assert!(node.started_at.is_none());
        assert!(node.last_heartbeat_at.is_none());

        // 2. Transition to Running
        dag.update_status("A", TaskStatus::Running);
        let node = dag.nodes.get("A").unwrap();
        assert_eq!(node.status, TaskStatus::Running);
        assert!(node.started_at.is_some());
        assert!(node.last_heartbeat_at.is_some());
        
        let start_time = node.started_at.unwrap();

        // 3. Update Heartbeat
        std::thread::sleep(std::time::Duration::from_millis(10));
        let now = chrono::Utc::now().timestamp_millis();
        dag.update_heartbeat("A", now);
        
        let node = dag.nodes.get("A").unwrap();
        assert!(node.last_heartbeat_at.unwrap() >= now);
        assert_eq!(node.started_at.unwrap(), start_time); // Start time shouldn't change

        // 4. Complete
        dag.update_status("A", TaskStatus::Completed);
        let node = dag.nodes.get("A").unwrap();
        assert_eq!(node.status, TaskStatus::Completed);
        // Timestamps remain
        assert!(node.started_at.is_some());
    }
}
