use crate::error::Result;
use crate::state::{AppState, ServiceOperationalStatus};
use common_contracts::messages::FetchCompanyDataCommand;
use common_contracts::workflow_types::WorkflowTaskCommand; // Import
use common_contracts::observability::ObservabilityTaskStatus;
use common_contracts::subjects::NatsSubject;
use common_contracts::ack::TaskAcknowledgement;
use futures_util::StreamExt;
use tracing::{error, info, warn};

use std::sync::Arc;
use tokio::sync::mpsc;
use chrono::Utc;
use std::time::Duration;

pub async fn run(state: AppState) -> Result<()> {
    info!("Starting NATS message consumer...");

    loop {
        // Always connect, regardless of Degraded status
        info!("Connecting to NATS...");
        match async_nats::connect(&state.config.nats_addr).await {
            Ok(client) => {
                info!("Successfully connected to NATS.");
                // Use try_join or spawn multiple subscribers
                let s1 = subscribe_legacy(state.clone(), client.clone());
                let s2 = subscribe_workflow(state.clone(), client.clone());
                
                tokio::select! {
                    res = s1 => if let Err(e) = res { error!("Legacy subscriber error: {}", e); },
                    res = s2 => if let Err(e) = res { error!("Workflow subscriber error: {}", e); },
                }
                
                // If any subscriber exits, wait and retry
            }
            Err(e) => {
                error!("Failed to connect to NATS: {}. Retrying in 10s...", e);
            }
        }
        tokio::time::sleep(Duration::from_secs(10)).await;
    }
}

async fn subscribe_workflow(state: AppState, client: async_nats::Client) -> Result<()> {
    let subject = "workflow.cmd.provider.tushare".to_string();
    let mut subscriber = client.subscribe(subject.clone()).await?;
    info!("Workflow Consumer started on '{}'", subject);

    while let Some(message) = subscriber.next().await {
        // Check Status (Handshake)
        let current_status = state.status.read().await.clone();
        
        // If Degraded, Reject immediately
        if let ServiceOperationalStatus::Degraded { reason } = current_status {
            warn!("Rejecting task due to degraded state: {}", reason);
            if let Some(reply_to) = message.reply {
                 let ack = TaskAcknowledgement::Rejected { reason };
                 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;
        }
        
        // If Active, Accept
        if let Some(reply_to) = message.reply.clone() {
             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 = state.clone();
        let client = client.clone();
        
        tokio::spawn(async move {
            match serde_json::from_slice::<WorkflowTaskCommand>(&message.payload) {
                Ok(cmd) => {
                    // TODO: Implement Heartbeat inside handle_workflow_command or wrapper
                    if let Err(e) = crate::generic_worker::handle_workflow_command(state, client, cmd).await {
                        error!("Generic worker handler failed: {}", e);
                    }
                },
                Err(e) => error!("Failed to parse WorkflowTaskCommand: {}", e),
            }
        });
    }
    Ok(())
}

async fn subscribe_legacy(
    state: AppState,
    client: async_nats::Client,
) -> Result<()> {
    let subject = NatsSubject::DataFetchCommands.to_string();
    let mut subscriber = client.subscribe(subject.clone()).await?;
    info!(
        "Legacy Consumer started, waiting for messages on subject '{}'",
        subject
    );

    while let Some(message) = subscriber.next().await {
        let current_status = state.status.read().await.clone();
        if matches!(current_status, ServiceOperationalStatus::Degraded {..}) {
            warn!("Service became degraded. Disconnecting from NATS and pausing consumption.");
            subscriber.unsubscribe().await?;
            return Ok(());
        }

        info!("Received NATS message.");
        let state_for_closure = Arc::new(state.clone());

        tokio::spawn(async move {
            if let Err(e) = serde_json::from_slice::<FetchCompanyDataCommand>(&message.payload) {
                error!("Failed to deserialize message: {}", e);
                warn!("Received non-json message: {:?}", message.payload);
                return;
            }

            let command = match serde_json::from_slice::<FetchCompanyDataCommand>(&message.payload) {
                Ok(c) => c,
                Err(e) => {
                    error!("Failed to deserialize message: {}", e);
                    return;
                }
            };
            info!("Received data fetch command for symbol: {}", command.symbol);

            // Tushare is for the Chinese market ("CN")
            if command.market.to_uppercase() != "CN" {
                info!(
                    "Skipping command for symbol '{}' as its market ('{}') is not 'CN'.",
                    command.symbol, command.market
                );
                return;
            }

            let (tx, rx) = mpsc::channel(1);
            let task_id = command.request_id;
            // Initialize deterministic progress entry
            state_for_closure.tasks.insert(task_id, common_contracts::observability::TaskProgress {
                request_id: task_id,
                task_name: format!("tushare:{}", command.symbol),
                status: ObservabilityTaskStatus::Queued,
                progress_percent: 0,
                details: "Command received".to_string(),
                started_at: Utc::now(),
            });

            // Spawn the workflow in a separate task
            let workflow_state = state_for_closure.clone();
            tokio::spawn(async move {
                let workflow_state_for_error = workflow_state.clone();
                let result = crate::worker::run_tushare_workflow(workflow_state, command, tx).await;
                if let Err(e) = result {
                    error!(
                        "Error executing Tushare workflow for task {}: {:?}",
                        task_id, e
                    );
                    // Update task to failed status
                    if let Some(mut task) = workflow_state_for_error.tasks.get_mut(&task_id) {
                        task.status = ObservabilityTaskStatus::Failed;
                        task.details = format!("Workflow failed: {}", e);
                    }
                }
            });

            // Spawn a separate task to clean up the task entry after completion or timeout
            let cleanup_state = state_for_closure.clone();
            tokio::spawn(async move {
                let mut rx = rx;
                match rx.recv().await {
                    Some(_) => {
                        info!("Task {} completed successfully, removing from map in 5 minutes.", task_id);
                        tokio::time::sleep(Duration::from_secs(300)).await;
                        cleanup_state.tasks.remove(&task_id);
                    }
                    None => {
                        // Check if the task is in Failed state (set by the worker error handler)
                        let is_failed = cleanup_state.tasks.get(&task_id)
                            .map(|t| t.status == ObservabilityTaskStatus::Failed)
                            .unwrap_or(false);

                        if is_failed {
                            info!("Task {} failed (handled), removing from map in 5 minutes.", task_id);
                            tokio::time::sleep(Duration::from_secs(300)).await;
                        } else {
                            warn!(
                                "Task {} completion signal not received (and not marked failed), removing after timeout.",
                                task_id
                            );
                        }
                        cleanup_state.tasks.remove(&task_id);
                    }
                }
            });
        });
    }
    Ok(())
}
