use crate::error::Result;
use crate::state::AppState;
use common_contracts::messages::FetchCompanyDataCommand;
use common_contracts::workflow_types::WorkflowTaskCommand;
use common_contracts::subjects::NatsSubject;
use futures_util::StreamExt;
use tracing::{error, info};

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

    let client = async_nats::connect(&state.config.nats_addr).await?;
    info!("Connected to NATS.");

    let legacy = subscribe_legacy(state.clone(), client.clone());
    let workflow = subscribe_workflow(state.clone(), client.clone());

    tokio::try_join!(legacy, workflow)?;

    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 {
        info!("Received Legacy NATS message.");
        let state_clone = state.clone();
        let publisher_clone = client.clone();

        tokio::spawn(async move {
            match serde_json::from_slice::<FetchCompanyDataCommand>(&message.payload) {
                Ok(command) => {
                    let request_id = command.request_id;
                    info!("Deserialized command for symbol: {}", command.symbol);
                    if let Err(e) =
                        crate::worker::handle_fetch_command(state_clone.clone(), command, publisher_clone)
                            .await
                    {
                        error!("Error handling fetch command: {:?}", e);
                        if let Some(mut task) = state_clone.tasks.get_mut(&request_id) {
                            task.status = common_contracts::observability::ObservabilityTaskStatus::Failed;
                            task.details = format!("Worker failed: {}", e);
                        }
                    }
                }
                Err(e) => {
                    error!("Failed to deserialize message: {}", e);
                }
            }
        });
    }
    Ok(())
}

use common_contracts::ack::TaskAcknowledgement;

async fn subscribe_workflow(state: AppState, client: async_nats::Client) -> Result<()> {
    let routing_key = "provider.yfinance".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.");
        
        // --- ACKNOWLEDGEMENT HANDSHAKE ---
        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_clone = state.clone();
        let client_clone = client.clone();

        tokio::spawn(async move {
            match serde_json::from_slice::<WorkflowTaskCommand>(&message.payload) {
                Ok(cmd) => {
                    if let Err(e) = crate::generic_worker::handle_workflow_command(state_clone, client_clone, cmd).await {
                        error!("Error handling workflow command: {:?}", e);
                    }
                }
                Err(e) => {
                    error!("Failed to deserialize workflow message: {}", e);
                }
            }
        });
    }
    Ok(())
}
