use std::sync::Arc;
use std::time::Duration;
use tracing::{info, warn, error};
use crate::state::AppState;
use common_contracts::workflow_types::{WorkflowTaskEvent, TaskStatus, TaskResult};
use common_contracts::subjects::SubjectMessage;

// Configuration
const TASK_MAX_DURATION_SEC: i64 = 300; // 5 minutes absolute timeout
const HEARTBEAT_TIMEOUT_SEC: i64 = 60;  // 60 seconds without heartbeat

pub async fn run(state: Arc<AppState>, nats: async_nats::Client) {
    info!("Task Monitor (Watchdog) started.");
    
    let mut interval = tokio::time::interval(Duration::from_secs(1));
    
    loop {
        interval.tick().await;
        
        let workflow_ids: Vec<uuid::Uuid> = state.workflows.iter().map(|r| *r.key()).collect();
        
        for req_id in workflow_ids {
            let mut events_to_publish = Vec::new();

            if let Some(dag_arc) = state.workflows.get(&req_id) {
                let dag = dag_arc.lock().await; 
                
                if dag.is_workflow_finished() {
                     continue;
                }
                
                let now = chrono::Utc::now().timestamp_millis();
                
                for (task_id, node) in dag.nodes.iter() {
                    if node.status == TaskStatus::Running {
                        let mut failure_reason = None;

                        // 1. Check Absolute Timeout
                        if let Some(start) = node.started_at {
                             if (now - start) > TASK_MAX_DURATION_SEC * 1000 {
                                 failure_reason = Some("Task execution timed out (5m limit)".to_string());
                             }
                        }
                        
                        // 2. Check Heartbeat
                        if failure_reason.is_none() {
                            if let Some(hb) = node.last_heartbeat_at {
                                if (now - hb) > HEARTBEAT_TIMEOUT_SEC * 1000 {
                                    failure_reason = Some("Task heartbeat lost (Zombie worker)".to_string());
                                }
                            } else if let Some(start) = node.started_at {
                                 // If started but no heartbeat yet
                                 if (now - start) > HEARTBEAT_TIMEOUT_SEC * 1000 {
                                      failure_reason = Some("Task unresponsive (No initial heartbeat)".to_string());
                                 }
                            }
                        }
                        
                        if let Some(reason) = failure_reason {
                             warn!("Watchdog detected failure for task {} in workflow {}: {}", task_id, req_id, reason);
                             
                             events_to_publish.push(WorkflowTaskEvent {
                                 request_id: req_id,
                                 task_id: task_id.clone(),
                                 status: TaskStatus::Failed,
                                 result: Some(TaskResult {
                                     new_commit: None,
                                     error: Some(reason),
                                     summary: None,
                                 }),
                             });
                        }
                    }
                }
            }
            
            // Publish events outside the lock
            for evt in events_to_publish {
                let subject = evt.subject().to_string();
                if let Ok(payload) = serde_json::to_vec(&evt) {
                    if let Err(e) = nats.publish(subject, payload.into()).await {
                         error!("Failed to publish watchdog event: {}", e);
                    }
                }
            }
        }
    }
}

