use anyhow::Result;
use tracing::{info, error};
use common_contracts::workflow_types::WorkflowTaskCommand;
use common_contracts::subjects::NatsSubject;
use common_contracts::ack::TaskAcknowledgement;
use crate::state::AppState;
use futures_util::StreamExt;
use std::sync::Arc;
use common_contracts::workflow_runner::WorkflowNodeRunner;
use crate::workflow_adapter::MockNode;

pub async fn run_consumer(state: AppState) -> Result<()> {
    info!("Starting Mock Provider message consumer...");
    let client = async_nats::connect(&state.config.nats_addr).await?;
    
    let routing_key = "provider.mock".to_string();
    let subject = NatsSubject::WorkflowCommand(routing_key).to_string();
    let mut subscriber = client.subscribe(subject.clone()).await?;

    info!("Workflow Consumer started on '{}'", subject);

    while let Some(message) = subscriber.next().await {
        info!("Received Workflow NATS message.");
        
        // 1. Parse Command eagerly to check config
        let cmd = match serde_json::from_slice::<WorkflowTaskCommand>(&message.payload) {
            Ok(c) => c,
            Err(e) => {
                error!("Failed to parse WorkflowTaskCommand: {}", e);
                continue;
            }
        };

        // 2. Check Simulation Mode
        let mode_raw = cmd.config.get("simulation_mode").and_then(|v| v.as_str()).unwrap_or("normal");
        let mode = mode_raw.to_lowercase();
        info!("Processing task {} with mode: {}", cmd.task_id, mode);
        
        if mode == "timeout_ack" {
            info!("Simulating Timeout (No ACK) for task {}", cmd.task_id);
            continue; // Skip processing
        }

        if mode == "reject" {
            info!("Simulating Rejection for task {}", cmd.task_id);
            if let Some(reply_to) = message.reply {
                 let ack = TaskAcknowledgement::Rejected { reason: "Simulated Rejection".into() };
                 if let Ok(payload) = serde_json::to_vec(&ack) {
                    if let Err(e) = client.publish(reply_to, payload.into()).await {
                        error!("Failed to send Rejection ACK: {}", e);
                    }
                 }
            }
            continue;
        }

        // 3. Normal / Crash / Hang Mode -> Send Accepted
        if let Some(reply_to) = message.reply {
             let ack = TaskAcknowledgement::Accepted;
             if let Ok(payload) = serde_json::to_vec(&ack) {
                if let Err(e) = client.publish(reply_to, payload.into()).await {
                    error!("Failed to send Acceptance ACK: {}", e);
                }
             }
        }

        let state_clone = state.clone();
        let client_clone = client.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_workflow_command(state_clone, client_clone, cmd).await {
                error!("Error handling workflow command: {:?}", e);
            }
        });
    }
    Ok(())
}

async fn handle_workflow_command(state: AppState, nats: async_nats::Client, cmd: WorkflowTaskCommand) -> Result<()> {
    let node = Arc::new(MockNode::new(state));
    let runner = WorkflowNodeRunner::new(nats);
    runner.run(node, cmd).await
}
