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, WorkflowEvent};
use common_contracts::workflow_types::WorkflowTaskEvent;
use common_contracts::subjects::NatsSubject;
use crate::workflow::WorkflowEngine;
use uuid::Uuid;

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

    // Topic 1: Workflow Commands (Start)
    let mut start_sub = nats.subscribe(NatsSubject::WorkflowCommandStart.to_string()).await?;
    
    // Topic 1b: Workflow Commands (Sync State)
    let mut sync_sub = nats.subscribe(NatsSubject::WorkflowCommandSyncState.to_string()).await?;

    // Topic 2: Workflow Task Events (Generic)
    let mut task_sub = nats.subscribe(NatsSubject::WorkflowEventTaskCompleted.to_string()).await?;

    // Topic 3: All Workflow Events (for capturing logs/stream)
    // events.workflow.> matches events.workflow.{req_id}
    let mut events_sub = nats.subscribe("events.workflow.>".to_string()).await?;

    let engine = Arc::new(WorkflowEngine::new(state.clone(), nats.clone()));

    // --- Task 1: Start Workflow ---
    let engine1 = engine.clone();
    tokio::spawn(async move {
        while let Some(msg) = start_sub.next().await {
            if let Ok(cmd) = serde_json::from_slice::<StartWorkflowCommand>(&msg.payload) {
                info!("Received StartWorkflow: {:?}", cmd);
                let engine_inner = engine1.clone();
                tokio::spawn(async move {
                    // Ensure handle_start_workflow is robust against Send constraints by wrapping if necessary
                    if let Err(e) = engine_inner.handle_start_workflow(cmd).await {
                        error!("Failed to handle StartWorkflow: {}", e);
                    }
                });
            } else {
                error!("Failed to parse StartWorkflowCommand");
            }
        }
    });

    // --- Task 1b: Sync State ---
    let engine_sync = engine.clone();
    tokio::spawn(async move {
        while let Some(msg) = sync_sub.next().await {
            if let Ok(cmd) = serde_json::from_slice::<SyncStateCommand>(&msg.payload) {
                info!("Received SyncStateCommand: request_id={}", cmd.request_id);
                let engine_inner = engine_sync.clone();
                tokio::spawn(async move {
                    if let Err(e) = engine_inner.handle_sync_state(cmd).await {
                        error!("Failed to handle SyncStateCommand: {}", e);
                    }
                });
            } else {
                error!("Failed to parse SyncStateCommand");
            }
        }
    });

    // --- Task 2: Task Completed Events ---
    let engine2 = engine.clone();
    tokio::spawn(async move {
        while let Some(msg) = task_sub.next().await {
            if let Ok(evt) = serde_json::from_slice::<WorkflowTaskEvent>(&msg.payload) {
                info!("Received TaskCompleted: task_id={}", evt.task_id);
                let engine_inner = engine2.clone();
                tokio::spawn(async move {
                    if let Err(e) = engine_inner.handle_task_completed(evt).await {
                        error!("Failed to handle TaskCompleted: {}", e);
                    }
                });
            } else {
                error!("Failed to parse WorkflowTaskEvent");
            }
        }
    });

    // --- Task 3: Workflow Events Capture (Logs & Stream) ---
    let engine_events = engine.clone();
    tokio::spawn(async move {
        while let Some(msg) = events_sub.next().await {
             // 1. Extract Request ID from Subject
            let subject_str = msg.subject.to_string();
            let parts: Vec<&str> = subject_str.split('.').collect();
            // Expected: events.workflow.{uuid}
            
            let req_id = if parts.len() >= 3 {
                match Uuid::parse_str(parts[2]) {
                    Ok(id) => id,
                    Err(_) => continue, // Ignore malformed subjects or non-UUID
                }
            } else {
                continue;
            };

            // 2. Parse Event Payload
            if let Ok(evt) = serde_json::from_slice::<WorkflowEvent>(&msg.payload) {
                let engine_inner = engine_events.clone();
                
                // 3. Dispatch based on event type
                match evt {
                    WorkflowEvent::TaskStreamUpdate { task_id, content_delta, .. } => {
                         tokio::spawn(async move {
                            let _ = engine_inner.handle_task_stream_update(task_id, content_delta, req_id).await;
                         });
                    },
                    WorkflowEvent::TaskLog { task_id, message, level, timestamp } => {
                         // Format log consistent with frontend expectations or raw?
                         // Frontend uses: `[${time}] [${p.level}] ${p.message}`
                         use chrono::{Utc, TimeZone};
                         let dt = Utc.timestamp_millis_opt(timestamp).unwrap();
                         let time_str = dt.format("%H:%M:%S").to_string();
                         let formatted_log = format!("[{}] [{}] {}", time_str, level, message);
                         
                         tokio::spawn(async move {
                            let _ = engine_inner.handle_task_log(task_id, formatted_log, req_id).await;
                         });
                    },
                    // Ignore others
                    _ => {}
                }
            }
        }
    });

    Ok(())
}
