use std::sync::Arc;
use common_contracts::workflow_types::{
    WorkflowTaskCommand, WorkflowTaskEvent, TaskStatus, StorageConfig
};
use common_contracts::messages::{
    StartWorkflowCommand, TaskType
};
use common_contracts::subjects::SubjectMessage;
use common_contracts::symbol_utils::CanonicalSymbol;
use tracing::{info, warn, error};
use anyhow::Result;
use serde_json::json;
use tokio::sync::Mutex;

use crate::dag_scheduler::DagScheduler;
use crate::state::AppState;
use workflow_context::{Vgcs, ContextStore}; // Added ContextStore

pub struct WorkflowEngine {
    state: Arc<AppState>,
    nats: async_nats::Client,
}

impl WorkflowEngine {
    pub fn new(state: Arc<AppState>, nats: async_nats::Client) -> Self {
        Self { state, nats }
    }

    pub async fn handle_start_workflow(&self, cmd: StartWorkflowCommand) -> Result<()> {
        let req_id = cmd.request_id;
        info!("Starting workflow {}", req_id);

        // 1. Init VGCS Repo
        self.state.vgcs.init_repo(&req_id.to_string())?;

        // 2. Create Scheduler
        // Initial commit is empty for a fresh workflow
        let mut dag = DagScheduler::new(req_id, String::new());

        // 3. Build DAG (Simplified Hardcoded logic for now, matching old build_dag)
        // In a real scenario, we fetch Template from DB/Service using cmd.template_id
        self.build_dag(&mut dag, &cmd.template_id, &cmd.market, &cmd.symbol);

        // 4. Save State
        self.state.workflows.insert(req_id, Arc::new(Mutex::new(dag.clone())));

        // 5. Trigger Initial Tasks
        let initial_tasks = dag.get_initial_tasks();
        
        // Lock the DAG for updates
        let dag_arc = self.state.workflows.get(&req_id).unwrap().clone();
        let mut dag_guard = dag_arc.lock().await;

        for task_id in initial_tasks {
            self.dispatch_task(&mut dag_guard, &task_id, &self.state.vgcs).await?;
        }

        Ok(())
    }

    pub async fn handle_task_completed(&self, evt: WorkflowTaskEvent) -> Result<()> {
        let req_id = evt.request_id;
        
        let dag_arc = match self.state.workflows.get(&req_id) {
            Some(d) => d.clone(),
            None => {
                warn!("Received event for unknown workflow {}", req_id);
                return Ok(());
            }
        };

        let mut dag = dag_arc.lock().await;

        // 1. Update Status & Record Commit
        dag.update_status(&evt.task_id, evt.status);
        
        if let Some(result) = evt.result {
            if let Some(commit) = result.new_commit {
                info!("Task {} produced commit {}", evt.task_id, commit);
                dag.record_result(&evt.task_id, Some(commit));
            }
            if let Some(err) = result.error {
                warn!("Task {} failed with error: {}", evt.task_id, err);
            }
        }

        // 2. Check for downstream tasks
        if evt.status == TaskStatus::Completed {
            let ready_tasks = dag.get_ready_downstream_tasks(&evt.task_id);
            for task_id in ready_tasks {
                if let Err(e) = self.dispatch_task(&mut dag, &task_id, &self.state.vgcs).await {
                    error!("Failed to dispatch task {}: {}", task_id, e);
                }
            }
        } else if evt.status == TaskStatus::Failed {
            // Handle failure propagation (skip downstream?)
            // For now, just log.
            error!("Task {} failed. Workflow might stall.", evt.task_id);
        }

        // 3. Check Workflow Completion (TODO)

        Ok(())
    }

    async fn dispatch_task(&self, dag: &mut DagScheduler, task_id: &str, vgcs: &Vgcs) -> Result<()> {
        // 1. Resolve Context (Merge if needed)
        let context = dag.resolve_context(task_id, vgcs)?;
        
        // 2. Update Status
        dag.update_status(task_id, TaskStatus::Scheduled);

        // 3. Construct Command
        let node = dag.nodes.get(task_id).ok_or_else(|| anyhow::anyhow!("Node not found"))?;
        
        let cmd = WorkflowTaskCommand {
            request_id: dag.request_id,
            task_id: task_id.to_string(),
            routing_key: node.routing_key.clone(),
            config: node.config.clone(),
            context,
            storage: StorageConfig {
                root_path: self.state.config.workflow_data_path.clone(),
            },
        };

        // 4. Publish
        let subject = cmd.subject().to_string(); // This uses the routing_key
        let payload = serde_json::to_vec(&cmd)?;
        
        info!("Dispatching task {} to subject {}", task_id, subject);
        self.nats.publish(subject, payload.into()).await?;

        Ok(())
    }

    // Helper to build DAG (Migrated from old workflow.rs)
    fn build_dag(&self, dag: &mut DagScheduler, template_id: &str, market: &str, symbol: &CanonicalSymbol) {
        // Logic copied/adapted from old WorkflowStateMachine::build_dag
        
        let mut providers = Vec::new();
        match market {
            "CN" => {
                providers.push("tushare");
                // providers.push("yfinance"); 
            },
            "US" => providers.push("yfinance"),
            _ => providers.push("yfinance"),
        }

        // 1. Data Fetch Nodes
        for p in &providers {
            let task_id = format!("fetch:{}", p);
            dag.add_node(
                task_id.clone(), 
                TaskType::DataFetch, 
                format!("provider.{}", p), // routing_key: workflow.cmd.provider.tushare
                json!({
                    "symbol": symbol.as_str(), // Simplification
                    "market": market
                })
            );
        }

        // 2. Analysis Node (Simplified)
        let report_task_id = "analysis:report";
        dag.add_node(
            report_task_id.to_string(),
            TaskType::Analysis,
            "analysis.report".to_string(), // routing_key: workflow.cmd.analysis.report
            json!({
                "template_id": template_id
            })
        );

        // 3. Edges
        for p in &providers {
            dag.add_dependency(&format!("fetch:{}", p), report_task_id);
        }
    }
}
