use std::sync::Arc;
use async_nats::Client;
use anyhow::Result;
use tracing::{info, error};
use futures::StreamExt;
use crate::state::AppState;
use common_contracts::messages::{
    StartWorkflowCommand, SyncStateCommand, 
    FinancialsPersistedEvent, DataFetchFailedEvent,
};
use common_contracts::subjects::NatsSubject;
use crate::persistence::{DashMapWorkflowRepository, NatsMessageBroker};
use crate::workflow::WorkflowEngine;

pub async fn run(state: Arc<AppState>, nats: Client) -> Result<()> {
    info!("Message Consumer started. Subscribing to topics...");

    let mut workflow_sub = nats.subscribe(NatsSubject::WorkflowCommandsWildcard.to_string()).await?;
    let mut data_sub = nats.subscribe(NatsSubject::DataEventsWildcard.to_string()).await?;
    let mut analysis_sub = nats.subscribe(NatsSubject::AnalysisEventsWildcard.to_string()).await?;

    // Prepare components
    let repo = DashMapWorkflowRepository::new(state.workflows.clone());
    let broker = NatsMessageBroker::new(nats.clone());

    // ... (Task 1 & 2 remain same) ...

    // --- Task 3: Analysis Events (Report Gen updates) ---
    let repo3 = repo.clone();
    let broker3 = broker.clone();
    tokio::spawn(async move {
        let engine = WorkflowEngine::new(Box::new(repo3), Box::new(broker3));
        while let Some(msg) = analysis_sub.next().await {
            match NatsSubject::try_from(msg.subject.as_str()) {
                Ok(NatsSubject::AnalysisReportGenerated) => {
                    if let Ok(evt) = serde_json::from_slice::<common_contracts::messages::ReportGeneratedEvent>(&msg.payload) {
                        info!("Received ReportGenerated: {:?}", evt);
                        if let Err(e) = engine.handle_report_generated(evt).await {
                            error!("Failed to handle ReportGenerated: {}", e);
                        }
                    }
                },
                Ok(NatsSubject::AnalysisReportFailed) => {
                    if let Ok(evt) = serde_json::from_slice::<common_contracts::messages::ReportFailedEvent>(&msg.payload) {
                        info!("Received ReportFailed: {:?}", evt);
                        if let Err(e) = engine.handle_report_failed(evt).await {
                            error!("Failed to handle ReportFailed: {}", e);
                        }
                    }
                },
                _ => {
                    // Ignore other subjects or log warning
                }
            }
        }
    });

    // --- Task 1: Workflow Commands (Start, Sync) ---
    let repo1 = repo.clone();
    let broker1 = broker.clone();
    tokio::spawn(async move {
        let engine = WorkflowEngine::new(Box::new(repo1), Box::new(broker1));
        while let Some(msg) = workflow_sub.next().await {
            match NatsSubject::try_from(msg.subject.as_str()) {
                Ok(NatsSubject::WorkflowCommandStart) => {
                    if let Ok(cmd) = serde_json::from_slice::<StartWorkflowCommand>(&msg.payload) {
                        info!("Received StartWorkflow: {:?}", cmd);
                        if let Err(e) = engine.handle_start_workflow(cmd).await {
                            error!("Failed to handle StartWorkflow: {}", e);
                        }
                    } else {
                        error!("Failed to parse StartWorkflowCommand");
                    }
                },
                Ok(NatsSubject::WorkflowCommandSyncState) => {
                    if let Ok(cmd) = serde_json::from_slice::<SyncStateCommand>(&msg.payload) {
                        info!("Received SyncState: {:?}", cmd);
                        // 关键修复：这里直接调用处理逻辑，如果还不行，我们需要检查 state 是否真的加载到了
                        // 或者在这里加更多的 log
                        match engine.handle_sync_state(cmd).await {
                            Ok(_) => info!("Successfully processed SyncState"),
                            Err(e) => error!("Failed to handle SyncState: {}", e),
                        }
                    }
                },
                _ => {}
            }
        }
    });

    // --- Task 2: Data Events (Progress updates) ---
    let repo2 = repo.clone();
    let broker2 = broker.clone();
    tokio::spawn(async move {
        let engine = WorkflowEngine::new(Box::new(repo2), Box::new(broker2));
        while let Some(msg) = data_sub.next().await {
            match NatsSubject::try_from(msg.subject.as_str()) {
                Ok(NatsSubject::DataFinancialsPersisted) => {
                    if let Ok(evt) = serde_json::from_slice::<FinancialsPersistedEvent>(&msg.payload) {
                        info!("Received DataFetched: {:?}", evt);
                        if let Err(e) = engine.handle_data_fetched(evt).await {
                            error!("Failed to handle DataFetched: {}", e);
                        }
                    }
                },
                Ok(NatsSubject::DataFetchFailed) => {
                    if let Ok(evt) = serde_json::from_slice::<DataFetchFailedEvent>(&msg.payload) {
                        info!("Received DataFailed: {:?}", evt);
                        if let Err(e) = engine.handle_data_failed(evt).await {
                            error!("Failed to handle DataFailed: {}", e);
                        }
                    }
                },
                _ => {}
            }
        }
    });

    Ok(())
}
