use std::sync::Arc;

use async_nats::jetstream::{self, consumer::pull, stream};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::mpsc;
use tracing::{error, info, warn};

use crate::registry::AgentRegistry;

/// Configuration for the NATS JetStream broker connection.
#[derive(Debug, Clone)]
pub struct NatsBrokerConfig {
    /// NATS server URL (e.g. `nats://localhost:4222`).
    pub url: String,
    /// JetStream stream name; created if it does not exist.
    pub stream_name: String,
    /// Subject for dispatching tasks into the stream.
    pub dispatch_subject: String,
    /// Subject consumed for task results (maps to `vapora.tasks.completed`).
    pub result_subject: String,
    /// Durable consumer name for result subscriptions.
    pub consumer_name: String,
}

impl Default for NatsBrokerConfig {
    fn default() -> Self {
        Self {
            url: "nats://localhost:4222".to_string(),
            stream_name: "VAPORA_TASKS".to_string(),
            dispatch_subject: "vapora.tasks.dispatch".to_string(),
            result_subject: crate::messages::subjects::TASKS_COMPLETED.to_string(),
            consumer_name: "vapora-nats-bridge".to_string(),
        }
    }
}

/// Task message dispatched to JetStream for agent processing.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMessage {
    pub task_id: String,
    pub title: String,
    pub description: String,
    pub role: String,
    pub context: String,
    pub priority: u32,
}

/// Task result received from the `vapora.tasks.completed` JetStream subject.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskResult {
    pub task_id: String,
    pub agent_id: String,
    pub result: String,
    pub success: bool,
    pub duration_ms: u64,
}

/// Agent status snapshot from the registry — never hardcoded.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStatusInfo {
    pub agent_id: String,
    pub role: String,
    pub available: bool,
    pub current_tasks: usize,
    pub max_tasks: usize,
}

#[derive(Debug, Error)]
pub enum NatsBridgeError {
    #[error("NATS connection failed: {0}")]
    Connection(String),

    #[error("JetStream stream error: {0}")]
    Stream(String),

    #[error("JetStream consumer error: {0}")]
    Consumer(String),

    #[error("Publish ack error: {0}")]
    Publish(String),

    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),
}

/// JetStream-backed bridge for durable task dispatch and result consumption.
///
/// Maintains a `VAPORA_TASKS` stream covering `vapora.tasks.dispatch` and
/// `vapora.tasks.completed`. Tasks are published with full JetStream acks;
/// results are consumed via a durable pull consumer so no messages are lost on
/// process restart.
pub struct NatsBridge {
    jetstream: jetstream::Context,
    config: NatsBrokerConfig,
    registry: Arc<AgentRegistry>,
}

impl NatsBridge {
    /// Connect to NATS and ensure the JetStream stream exists.
    ///
    /// The stream uses `WorkQueue` retention so each message is delivered to
    /// exactly one consumer and removed after ack.
    pub async fn connect(
        config: NatsBrokerConfig,
        registry: Arc<AgentRegistry>,
    ) -> Result<Self, NatsBridgeError> {
        let client = async_nats::connect(&config.url)
            .await
            .map_err(|e| NatsBridgeError::Connection(e.to_string()))?;

        let jetstream = jetstream::new(client);

        let stream_config = stream::Config {
            name: config.stream_name.clone(),
            subjects: vec![
                config.dispatch_subject.clone(),
                config.result_subject.clone(),
            ],
            retention: stream::RetentionPolicy::WorkQueue,
            storage: stream::StorageType::File,
            num_replicas: 1,
            ..Default::default()
        };

        jetstream
            .get_or_create_stream(stream_config)
            .await
            .map_err(|e| NatsBridgeError::Stream(e.to_string()))?;

        info!(
            url = %config.url,
            stream = %config.stream_name,
            dispatch = %config.dispatch_subject,
            result = %config.result_subject,
            "NatsBridge connected; JetStream stream ensured"
        );

        Ok(Self {
            jetstream,
            config,
            registry,
        })
    }

    /// Publish a task to JetStream with a synchronous ack.
    ///
    /// Returns the stream sequence number assigned to this message. The ack
    /// guarantees the message reached the NATS server and was persisted in the
    /// stream before returning.
    pub async fn submit_task(&self, task: TaskMessage) -> Result<u64, NatsBridgeError> {
        let payload = serde_json::to_vec(&task)?;

        let ack = self
            .jetstream
            .publish(self.config.dispatch_subject.clone(), payload.into())
            .await
            .map_err(|e| NatsBridgeError::Publish(e.to_string()))?
            .await
            .map_err(|e| NatsBridgeError::Publish(e.to_string()))?;

        info!(
            task_id = %task.task_id,
            sequence = ack.sequence,
            subject = %self.config.dispatch_subject,
            "Task dispatched to JetStream"
        );

        Ok(ack.sequence)
    }

    /// Start a background pull consumer for `vapora.tasks.completed`.
    ///
    /// Spawns a tokio task that continuously fetches from the durable consumer
    /// and sends deserialized `TaskResult` values into the returned channel.
    /// Each message is ack-ed after successful deserialization; failed
    /// deserialization sends a Nak so the message is redelivered.
    ///
    /// The consumer is durable — if the process restarts it will resume from
    /// the last unacked position.
    pub async fn subscribe_task_results(
        &self,
    ) -> Result<mpsc::Receiver<TaskResult>, NatsBridgeError> {
        let stream = self
            .jetstream
            .get_stream(&self.config.stream_name)
            .await
            .map_err(|e| NatsBridgeError::Stream(e.to_string()))?;

        let consumer: async_nats::jetstream::consumer::Consumer<pull::Config> = stream
            .get_or_create_consumer(
                &self.config.consumer_name,
                pull::Config {
                    durable_name: Some(self.config.consumer_name.clone()),
                    filter_subject: self.config.result_subject.clone(),
                    ..Default::default()
                },
            )
            .await
            .map_err(|e| NatsBridgeError::Consumer(e.to_string()))?;

        let (tx, rx) = mpsc::channel::<TaskResult>(256);

        tokio::spawn(run_result_consumer(consumer, tx));

        info!(
            consumer = %self.config.consumer_name,
            subject = %self.config.result_subject,
            "Task result pull consumer started"
        );

        Ok(rx)
    }

    /// List all registered agents from the registry.
    ///
    /// Never returns hardcoded data — always reflects the live registry state.
    pub fn list_agents(&self) -> Vec<AgentStatusInfo> {
        self.registry
            .list_all()
            .into_iter()
            .map(|agent| {
                let available = agent.can_accept_task();
                AgentStatusInfo {
                    agent_id: agent.id,
                    role: agent.role,
                    available,
                    current_tasks: agent.current_tasks as usize,
                    max_tasks: agent.max_concurrent_tasks as usize,
                }
            })
            .collect()
    }
}

/// Pull consumer loop: fetches JetStream messages, deserializes, acks, and
/// forwards `TaskResult` values to `tx`. Extracted to avoid excessive nesting
/// inside the `tokio::spawn` closure.
async fn run_result_consumer(
    consumer: async_nats::jetstream::consumer::Consumer<pull::Config>,
    tx: mpsc::Sender<TaskResult>,
) {
    let mut messages = match consumer.messages().await {
        Ok(m) => m,
        Err(e) => {
            error!(error = %e, "Failed to create JetStream message stream");
            return;
        }
    };

    while let Some(msg_result) = futures::StreamExt::next(&mut messages).await {
        let msg = match msg_result {
            Ok(m) => m,
            Err(e) => {
                error!(error = %e, "JetStream message error; stopping consumer");
                break;
            }
        };

        match serde_json::from_slice::<TaskResult>(&msg.payload) {
            Ok(result) => {
                if let Err(e) = msg.ack().await {
                    error!(error = %e, "Failed to ack JetStream message");
                    continue;
                }
                if tx.send(result).await.is_err() {
                    warn!("Result receiver dropped; stopping consumer loop");
                    break;
                }
            }
            Err(e) => {
                warn!(error = %e, "Failed to deserialize TaskResult; sending Nak");
                if let Err(nak_err) = msg
                    .ack_with(async_nats::jetstream::AckKind::Nak(None))
                    .await
                {
                    error!(error = %nak_err, "Failed to Nak message");
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::{AgentMetadata, AgentRegistry};

    fn make_registry_with_agent() -> Arc<AgentRegistry> {
        let reg = Arc::new(AgentRegistry::new(5));
        let agent = AgentMetadata::new(
            "developer".to_string(),
            "Dev 1".to_string(),
            "claude".to_string(),
            "claude-sonnet-4-6".to_string(),
            vec!["coding".to_string()],
        );
        reg.register_agent(agent).unwrap();
        reg
    }

    #[test]
    fn list_agents_reflects_registry() {
        let registry = make_registry_with_agent();
        // NatsBridge::list_agents requires a connected bridge, so test via a
        // registry snapshot directly to verify the mapping logic.
        let agents: Vec<AgentStatusInfo> = registry
            .list_all()
            .into_iter()
            .map(|agent| {
                let available = agent.can_accept_task();
                AgentStatusInfo {
                    agent_id: agent.id,
                    role: agent.role,
                    available,
                    current_tasks: agent.current_tasks as usize,
                    max_tasks: agent.max_concurrent_tasks as usize,
                }
            })
            .collect();

        assert_eq!(agents.len(), 1);
        assert_eq!(agents[0].role, "developer");
        assert!(agents[0].available);
        assert_eq!(agents[0].current_tasks, 0);
        assert_eq!(agents[0].max_tasks, 5);
    }

    #[test]
    fn task_message_roundtrip_serialization() {
        let msg = TaskMessage {
            task_id: "t-1".to_string(),
            title: "Fix bug".to_string(),
            description: "Critical regression in auth".to_string(),
            role: "developer".to_string(),
            context: "{}".to_string(),
            priority: 90,
        };
        let bytes = serde_json::to_vec(&msg).unwrap();
        let decoded: TaskMessage = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(decoded.task_id, msg.task_id);
        assert_eq!(decoded.priority, msg.priority);
    }

    #[test]
    fn task_result_roundtrip_serialization() {
        let result = TaskResult {
            task_id: "t-1".to_string(),
            agent_id: "agent-dev-1".to_string(),
            result: "Bug fixed in auth.rs:42".to_string(),
            success: true,
            duration_ms: 4200,
        };
        let bytes = serde_json::to_vec(&result).unwrap();
        let decoded: TaskResult = serde_json::from_slice(&bytes).unwrap();
        assert!(decoded.success);
        assert_eq!(decoded.duration_ms, 4200);
    }

    /// Integration test: requires NATS server at nats://localhost:4222.
    /// Run with: cargo test -p vapora-agents nats_bridge -- --ignored
    #[tokio::test]
    #[ignore]
    async fn submit_task_receives_ack() {
        let registry = make_registry_with_agent();
        let config = NatsBrokerConfig::default();

        let bridge = NatsBridge::connect(config, registry)
            .await
            .expect("NATS must be running at localhost:4222");

        let task = TaskMessage {
            task_id: "integration-t-1".to_string(),
            title: "Integration test task".to_string(),
            description: "Verify JetStream publish ack".to_string(),
            role: "developer".to_string(),
            context: "{}".to_string(),
            priority: 50,
        };

        let sequence = bridge.submit_task(task).await.unwrap();
        assert!(sequence > 0, "JetStream ack sequence must be > 0");
    }

    /// Integration test: requires NATS server at nats://localhost:4222.
    #[tokio::test]
    #[ignore]
    async fn subscribe_task_results_receives_published_result() {
        use tokio::time::{timeout, Duration};

        let registry = make_registry_with_agent();
        let config = NatsBrokerConfig::default();

        let bridge = NatsBridge::connect(config.clone(), registry)
            .await
            .expect("NATS must be running at localhost:4222");

        let mut rx = bridge.subscribe_task_results().await.unwrap();

        // Publish a result directly to the result subject via JetStream
        let result = TaskResult {
            task_id: "integration-t-2".to_string(),
            agent_id: "agent-dev-1".to_string(),
            result: "Done".to_string(),
            success: true,
            duration_ms: 100,
        };
        let payload = serde_json::to_vec(&result).unwrap();
        bridge
            .jetstream
            .publish(config.result_subject.clone(), payload.into())
            .await
            .unwrap()
            .await
            .unwrap();

        let received = timeout(Duration::from_secs(5), rx.recv())
            .await
            .expect("Timeout waiting for result")
            .expect("Channel closed");

        assert_eq!(received.task_id, "integration-t-2");
        assert!(received.success);
    }
}
