use std::sync::Arc;

use dashmap::DashMap;
use tokio::sync::{mpsc, oneshot};
use tracing::{error, info, warn};
use vapora_agents::{
    coordinator::AgentCoordinator,
    nats_bridge::{AgentStatusInfo, NatsBridge, TaskMessage, TaskResult},
};

use crate::{
    error::{A2aError, Result},
    metrics::{A2A_COORDINATOR_ASSIGNMENTS, A2A_NATS_MESSAGES, A2A_TASKS_TOTAL},
    protocol::{A2aErrorObj, A2aMessage, A2aMessagePart, A2aTask, A2aTaskResult, TaskState},
    task_manager::TaskManager,
};

pub struct CoordinatorBridge {
    coordinator: Arc<AgentCoordinator>,
    task_manager: Arc<TaskManager>,
    result_channels: Arc<DashMap<String, oneshot::Sender<A2aTaskResult>>>,
    nats_bridge: Option<Arc<NatsBridge>>,
}

impl CoordinatorBridge {
    pub fn new(
        coordinator: Arc<AgentCoordinator>,
        task_manager: Arc<TaskManager>,
        nats_bridge: Option<Arc<NatsBridge>>,
    ) -> Self {
        Self {
            coordinator,
            task_manager,
            result_channels: Arc::new(DashMap::new()),
            nats_bridge,
        }
    }

    /// Start the JetStream result consumer.
    ///
    /// Subscribes to `vapora.tasks.completed` via the durable pull consumer in
    /// `NatsBridge`. Each `TaskResult` received is persisted to SurrealDB and
    /// forwarded to any waiting oneshot channel. Replaces the previous core
    /// NATS subscriber approach with at-least-once JetStream delivery.
    pub async fn start_result_listener(&self) -> Result<()> {
        let Some(bridge) = &self.nats_bridge else {
            warn!("NatsBridge not configured; result listener disabled");
            return Ok(());
        };

        let rx = bridge
            .subscribe_task_results()
            .await
            .map_err(|e| A2aError::InternalError(format!("JetStream consumer error: {}", e)))?;

        let task_manager = self.task_manager.clone();
        let result_channels = self.result_channels.clone();

        tokio::spawn(run_result_consumer(rx, task_manager, result_channels));

        info!("A2A result listener started (JetStream pull consumer)");
        Ok(())
    }

    pub async fn dispatch(&self, a2a_task: A2aTask) -> Result<String> {
        let task_id = a2a_task.id.clone();

        let skill = a2a_task
            .metadata
            .get("skill")
            .and_then(|v| v.as_str())
            .unwrap_or("developer")
            .to_string();

        let task_text = a2a_task
            .message
            .parts
            .iter()
            .filter_map(|p| match p {
                A2aMessagePart::Text(t) => Some(t.clone()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("\n");

        if task_text.is_empty() {
            return Err(A2aError::InternalError(
                "No text in task message".to_string(),
            ));
        }

        let parts: Vec<&str> = task_text.splitn(2, '\n').collect();
        let title = parts[0].to_string();
        let description = parts.get(1).unwrap_or(&"").to_string();

        // Persist task (status: waiting)
        self.task_manager.create(a2a_task.clone()).await?;
        A2A_TASKS_TOTAL.with_label_values(&["waiting"]).inc();

        self.task_manager
            .update_state(&task_id, TaskState::Working)
            .await?;
        A2A_TASKS_TOTAL.with_label_values(&["working"]).inc();

        // Assign via AgentCoordinator (learning-based selection + core NATS publish)
        match self
            .coordinator
            .assign_task(
                &skill,
                title.clone(),
                description.clone(),
                serde_json::json!({}).to_string(),
                50,
            )
            .await
        {
            Ok(_) => {
                A2A_COORDINATOR_ASSIGNMENTS
                    .with_label_values(&[skill.as_str(), "success"])
                    .inc();
                info!(task_id = %task_id, "Task dispatched via coordinator");
            }
            Err(e) => {
                A2A_COORDINATOR_ASSIGNMENTS
                    .with_label_values(&[skill.as_str(), "error"])
                    .inc();
                return Err(A2aError::CoordinatorError(e.to_string()));
            }
        }

        // Publish to JetStream for durable delivery (if NatsBridge is configured)
        if let Some(bridge) = &self.nats_bridge {
            let task_msg = TaskMessage {
                task_id: task_id.clone(),
                title,
                description,
                role: skill.clone(),
                context: serde_json::json!({}).to_string(),
                priority: 50,
            };

            match bridge.submit_task(task_msg).await {
                Ok(seq) => {
                    info!(
                        task_id = %task_id,
                        jetstream_seq = seq,
                        "Task published to JetStream"
                    );
                }
                Err(e) => {
                    // Non-fatal: core NATS dispatch already succeeded.
                    // Log and continue — coordinator assignment was the primary path.
                    warn!(
                        error = %e,
                        task_id = %task_id,
                        "JetStream publish failed (task still dispatched via coordinator)"
                    );
                }
            }
        }

        Ok(task_id)
    }

    pub async fn get_task(&self, id: &str) -> Result<crate::protocol::A2aTaskStatus> {
        self.task_manager.get(id).await
    }

    /// List all registered agents from the NatsBridge registry.
    /// Returns empty vec if NatsBridge is not configured.
    pub fn list_agents(&self) -> Vec<AgentStatusInfo> {
        self.nats_bridge
            .as_ref()
            .map(|b| b.list_agents())
            .unwrap_or_default()
    }
}

async fn run_result_consumer(
    mut rx: mpsc::Receiver<TaskResult>,
    task_manager: Arc<TaskManager>,
    result_channels: Arc<DashMap<String, oneshot::Sender<A2aTaskResult>>>,
) {
    while let Some(task_result) = rx.recv().await {
        let task_id = task_result.task_id.clone();

        if task_result.success {
            let a2a_result = A2aTaskResult {
                message: A2aMessage {
                    role: "assistant".to_string(),
                    parts: vec![A2aMessagePart::Text(task_result.result.clone())],
                },
                artifacts: None,
            };

            match task_manager.complete(&task_id, a2a_result.clone()).await {
                Err(e) => {
                    error!(error = %e, task_id = %task_id, "Failed to mark task completed");
                    A2A_NATS_MESSAGES
                        .with_label_values(&["completed", "db_error"])
                        .inc();
                }
                Ok(()) => {
                    A2A_NATS_MESSAGES
                        .with_label_values(&["completed", "success"])
                        .inc();
                    A2A_TASKS_TOTAL.with_label_values(&["completed"]).inc();
                }
            }

            if let Some((_, sender)) = result_channels.remove(&task_id) {
                if sender.send(a2a_result).is_err() {
                    warn!(task_id = %task_id, "Result receiver dropped");
                }
            }
        } else {
            let error_obj = A2aErrorObj {
                code: -1,
                message: task_result.result.clone(),
            };

            match task_manager.fail(&task_id, error_obj).await {
                Err(e) => {
                    error!(error = %e, task_id = %task_id, "Failed to mark task failed");
                    A2A_NATS_MESSAGES
                        .with_label_values(&["failed", "db_error"])
                        .inc();
                }
                Ok(()) => {
                    A2A_NATS_MESSAGES
                        .with_label_values(&["failed", "success"])
                        .inc();
                    A2A_TASKS_TOTAL.with_label_values(&["failed"]).inc();
                }
            }

            result_channels.remove(&task_id);
        }
    }

    warn!("NatsBridge result channel closed; listener stopped");
}
