use std::sync::Arc;

use common_contracts::messages::GenerateReportCommand;
use common_contracts::workflow_types::{WorkflowTaskCommand, WorkflowTaskEvent, TaskStatus};
use common_contracts::symbol_utils::{CanonicalSymbol, Market};
// use common_contracts::subjects::SubjectMessage; // Removed unused import
use futures::StreamExt;
use tracing::{error, info};

use crate::{state::AppState, worker::run_report_generation_workflow};

const OLD_SUBJECT_NAME: &str = "analysis.commands.generate_report";
const WORKFLOW_CMD_SUBJECT: &str = "workflow.cmd.analysis.report";

pub async fn subscribe_to_commands(
    app_state: AppState,
    nats_client: async_nats::Client,
) -> Result<(), anyhow::Error> {
    let mut old_subscriber = nats_client.subscribe(OLD_SUBJECT_NAME.to_string()).await?;
    let mut workflow_subscriber = nats_client.subscribe(WORKFLOW_CMD_SUBJECT.to_string()).await?;

    info!(
        "Consumer started, waiting for commands on subjects '{}' and '{}'",
        OLD_SUBJECT_NAME, WORKFLOW_CMD_SUBJECT
    );

    // Spawn task for old command style
    let state_clone1 = app_state.clone();
    tokio::spawn(async move {
        while let Some(message) = old_subscriber.next().await {
            info!("Received legacy NATS command. Payload size: {} bytes", message.payload.len());
            let state = state_clone1.clone();
            tokio::spawn(async move {
                match serde_json::from_slice::<GenerateReportCommand>(&message.payload) {
                    Ok(command) => {
                        info!(
                            "Command triggered analysis for symbol: {}, template: {}, request_id: {}",
                            command.symbol, command.template_id, command.request_id
                        );
                        if let Err(e) = run_report_generation_workflow(Arc::new(state), command).await {
                            error!("Error running report generation workflow: {:?}", e);
                        }
                    }
                    Err(e) => {
                        error!("Failed to deserialize GenerateReportCommand: {}. Payload: {:?}", e, message.payload);
                    }
                }
            });
        }
    });

    // Spawn task for new WorkflowTaskCommand
    let state_clone2 = app_state.clone();
    let nats_client2 = nats_client.clone();
    tokio::spawn(async move {
        while let Some(message) = workflow_subscriber.next().await {
            info!("Received WorkflowTaskCommand. Payload size: {} bytes", message.payload.len());
            let state = state_clone2.clone();
            let nats = nats_client2.clone();
            
            tokio::spawn(async move {
                match serde_json::from_slice::<WorkflowTaskCommand>(&message.payload) {
                    Ok(task_cmd) => {
                        info!("Received WorkflowTaskCommand for task_id: {}", task_cmd.task_id);

                        // --- 0. Immediate Acknowledgement ---
                        if let Some(reply_subject) = message.reply.clone() {
                            let ack = common_contracts::ack::TaskAcknowledgement::Accepted;
                            if let Ok(payload) = serde_json::to_vec(&ack) {
                                if let Err(e) = nats.publish(reply_subject, payload.into()).await {
                                    error!("Failed to send ACK for task {}: {}", task_cmd.task_id, e);
                                } else {
                                    info!("ACK sent for task {}", task_cmd.task_id);
                                }
                            }
                        } else {
                            // This should only happen for fire-and-forget dispatch, which orchestrator doesn't use
                            // but logging it is good.
                            tracing::warn!("No reply subject for task {}, cannot send ACK.", task_cmd.task_id);
                        }

                        // 1. Extract params from config
                        let symbol_str = task_cmd.config.get("symbol").and_then(|v| v.as_str());
                        let market_str = task_cmd.config.get("market").and_then(|v| v.as_str());
                        let template_id = task_cmd.config.get("template_id").and_then(|v| v.as_str());
                        let module_id = task_cmd.config.get("module_id").and_then(|v| v.as_str());

                        if let (Some(s), Some(m), Some(t)) = (symbol_str, market_str, template_id) {
                            let market = Market::from(m);
                            let symbol = CanonicalSymbol::new(s, &market);
                            
                            let commit_hash = task_cmd.context.base_commit.clone();
                            let input_bindings = task_cmd.config.get("input_bindings").and_then(|v| serde_json::from_value(v.clone()).ok());
                            let output_path = task_cmd.config.get("output_path").and_then(|v| v.as_str()).map(|s| s.to_string());
                            
                            let mut llm_config = None;
                            let mut analysis_prompt = None;
                            
                            if let Some(mc_val) = task_cmd.config.get("_module_config") {
                                 if let Ok(mc) = serde_json::from_value::<common_contracts::config_models::AnalysisModuleConfig>(mc_val.clone()) {
                                     llm_config = mc.llm_config;
                                     analysis_prompt = Some(mc.analysis_prompt);
                                 }
                            }

                            let report_cmd = GenerateReportCommand {
                                request_id: task_cmd.request_id,
                                symbol: symbol.clone(),
                                template_id: t.to_string(),
                                task_id: Some(task_cmd.task_id.clone()),
                                module_id: module_id.map(|v| v.to_string()),
                                commit_hash,
                                input_bindings,
                                output_path,
                                llm_config,
                                analysis_prompt,
                            };

                            // 2. Send TaskStatus::Running
                            let running_evt = WorkflowTaskEvent {
                                request_id: task_cmd.request_id,
                                task_id: task_cmd.task_id.clone(),
                                status: TaskStatus::Running,
                                result: None,
                            };
                            
                            if let Ok(payload) = serde_json::to_vec(&running_evt) {
                                let subject = common_contracts::subjects::NatsSubject::WorkflowEventTaskCompleted.to_string();
                                if let Err(e) = nats.publish(subject, payload.into()).await {
                                     error!("Failed to publish TaskRunning event: {}", e);
                                }
                            }
                            
                            // 3. Run Logic
                            match run_report_generation_workflow(Arc::new(state), report_cmd).await {
                                Ok(new_commit_opt) => {
                                    // 4. Send TaskStatus::Completed
                                    let completed_evt = WorkflowTaskEvent {
                                        request_id: task_cmd.request_id,
                                        task_id: task_cmd.task_id.clone(),
                                        status: TaskStatus::Completed,
                                        result: Some(common_contracts::workflow_types::TaskResult {
                                            new_commit: new_commit_opt, // Pass the commit hash back
                                            error: None,
                                            summary: None,
                                        }),
                                    };
                                    if let Ok(payload) = serde_json::to_vec(&completed_evt) {
                                        let subject = common_contracts::subjects::NatsSubject::WorkflowEventTaskCompleted.to_string();
                                        if let Err(e) = nats.publish(subject, payload.into()).await {
                                             error!("Failed to publish TaskCompleted event: {}", e);
                                        }
                                    }
                                }
                                Err(e) => {
                                    error!("Workflow execution failed: {}", e);
                                    // 5. Send TaskStatus::Failed
                                    let failed_evt = WorkflowTaskEvent {
                                        request_id: task_cmd.request_id,
                                        task_id: task_cmd.task_id.clone(),
                                        status: TaskStatus::Failed,
                                        result: Some(common_contracts::workflow_types::TaskResult {
                                            new_commit: None,
                                            error: Some(e.to_string()),
                                            summary: None,
                                        }),
                                    };
                                    if let Ok(payload) = serde_json::to_vec(&failed_evt) {
                                        let subject = common_contracts::subjects::NatsSubject::WorkflowEventTaskCompleted.to_string();
                                        if let Err(e) = nats.publish(subject, payload.into()).await {
                                             error!("Failed to publish TaskFailed event: {}", e);
                                        }
                                    }
                                }
                            }

                        } else {
                            error!("Missing required config params in WorkflowTaskCommand: {:?}", task_cmd.config);
                            // Send Failed Event immediately
                             let failed_evt = WorkflowTaskEvent {
                                request_id: task_cmd.request_id,
                                task_id: task_cmd.task_id.clone(),
                                status: TaskStatus::Failed,
                                result: Some(common_contracts::workflow_types::TaskResult {
                                    new_commit: None,
                                    error: Some("Missing required config params (symbol, market, template_id)".to_string()),
                                    summary: None,
                                }),
                            };
                            if let Ok(payload) = serde_json::to_vec(&failed_evt) {
                                let subject = common_contracts::subjects::NatsSubject::WorkflowEventTaskCompleted.to_string();
                                let _ = nats.publish(subject, payload.into()).await;
                            }
                        }
                    }
                    Err(e) => {
                        error!("Failed to deserialize WorkflowTaskCommand: {}. Payload: {:?}", e, message.payload);
                    }
                }
            });
        }
    });

    Ok(())
}
