//! NATS JetStream bridge using `platform-nats` with NKey auth and signed messages.

use bytes::Bytes;
use platform_nats::{EventStream, NatsConfig};
use serde::{Deserialize, Serialize};
use std::time::Duration;

use crate::VaporaResult;

/// Broker-level configuration for `NatsBridge`.
///
/// Call `NatsBrokerConfig::to_platform_config()` to convert to the
/// `platform_nats::NatsConfig` expected by `EventStream::connect`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NatsBrokerConfig {
    /// NATS server URL, e.g. `nats://localhost:4222`
    pub url: String,
    /// JetStream stream name (created if absent)
    pub stream_name: String,
    /// Durable pull-consumer name (created if absent)
    pub consumer_name: String,
    /// Subject filter list, e.g. `["syntaxis.>"]`
    pub subjects: Vec<String>,
    /// ed25519 NKey seed for signing outbound messages (optional)
    pub nkey_seed: Option<String>,
    /// Public NKeys whose signatures are accepted (empty → accept all)
    pub trusted_nkeys: Vec<String>,
    /// Reject inbound messages that lack a valid NKey signature
    pub require_signed_messages: bool,
}

impl Default for NatsBrokerConfig {
    fn default() -> Self {
        Self {
            url: "nats://localhost:4222".to_string(),
            stream_name: "SYNTAXIS".to_string(),
            consumer_name: "syntaxis-consumer".to_string(),
            subjects: vec!["syntaxis.>".to_string()],
            nkey_seed: None,
            trusted_nkeys: vec![],
            require_signed_messages: false,
        }
    }
}

impl NatsBrokerConfig {
    fn to_platform_config(&self) -> NatsConfig {
        NatsConfig {
            url: self.url.clone(),
            stream_name: self.stream_name.clone(),
            consumer_name: self.consumer_name.clone(),
            subjects: self.subjects.clone(),
            nkey_seed: self.nkey_seed.clone(),
            trusted_nkeys: self.trusted_nkeys.clone(),
            require_signed_messages: self.require_signed_messages,
        }
    }
}

/// Event configuration matching the kogral `NatsEventConfig` schema.
///
/// Allows downstream systems to express NATS connectivity in a unified format
/// and convert to `NatsBrokerConfig` via `From`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NatsEventConfig {
    /// NATS server URL
    pub url: String,
    /// JetStream stream name
    pub stream: String,
    /// Durable pull-consumer name
    pub consumer: String,
    /// Subject filter list
    pub subjects: Vec<String>,
    /// NKey seed for signing outbound messages
    #[serde(default)]
    pub nkey_seed: Option<String>,
    /// Trusted public NKeys
    #[serde(default)]
    pub trusted_nkeys: Vec<String>,
    /// Reject inbound messages lacking a valid NKey signature
    #[serde(default)]
    pub require_signed_messages: bool,
}

impl From<NatsEventConfig> for NatsBrokerConfig {
    fn from(cfg: NatsEventConfig) -> Self {
        Self {
            url: cfg.url,
            stream_name: cfg.stream,
            consumer_name: cfg.consumer,
            subjects: cfg.subjects,
            nkey_seed: cfg.nkey_seed,
            trusted_nkeys: cfg.trusted_nkeys,
            require_signed_messages: cfg.require_signed_messages,
        }
    }
}

/// Task message published to `syntaxis.tasks.<agent_role>.submit`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMessage {
    /// Unique task identifier
    pub task_id: String,
    /// Agent role this task is targeted at
    pub agent_role: String,
    /// Short task title
    pub title: String,
    /// Full task description
    pub description: String,
    /// Priority in `1..=10` (higher = more urgent)
    pub priority: u32,
    /// Arbitrary structured task data
    pub data: serde_json::Value,
}

/// Task result consumed from `syntaxis.tasks.<task_id>.result`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskResult {
    /// Task identifier correlating to the originating `TaskMessage`
    pub task_id: String,
    /// Agent that completed (or failed) the task
    pub agent_role: String,
    /// `"completed"`, `"failed"`, or `"timeout"`
    pub status: String,
    /// Result payload produced by the agent
    pub result: serde_json::Value,
    /// Wall-clock execution time in milliseconds
    pub execution_time_ms: u64,
}

/// Agent heartbeat published to `syntaxis.agents.<role>.status`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStatus {
    /// Agent role identifier
    pub role: String,
    /// Whether the agent is currently healthy
    pub healthy: bool,
    /// Number of tasks currently in flight
    pub task_count: usize,
    /// Maximum concurrent task capacity
    pub capacity: usize,
    /// RFC 3339 timestamp of the last heartbeat
    pub last_heartbeat: String,
}

/// Production NATS JetStream bridge backed by `platform-nats::EventStream`.
pub struct NatsBridge {
    stream: EventStream,
    config: NatsBrokerConfig,
}

impl std::fmt::Debug for NatsBridge {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NatsBridge")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

impl NatsBridge {
    /// Connect to NATS and materialise the configured stream + consumer.
    pub async fn connect(config: NatsBrokerConfig) -> VaporaResult<Self> {
        let stream = EventStream::connect(&config.to_platform_config()).await?;
        tracing::info!(url = %config.url, stream = %config.stream_name, "NATS bridge connected");
        Ok(Self { stream, config })
    }

    /// Publish a task to `syntaxis.tasks.<agent_role>.submit`.
    ///
    /// Returns the `task_id` on success for correlation tracking.
    pub async fn submit_task(&self, task: &TaskMessage) -> VaporaResult<String> {
        let subject = format!("syntaxis.tasks.{}.submit", task.agent_role);
        let payload = Bytes::from(serde_json::to_vec(task)?);
        self.stream.publish(&subject, payload).await?;
        tracing::debug!(task_id = %task.task_id, subject = %subject, "task submitted");
        Ok(task.task_id.clone())
    }

    /// Pull up to `max_msgs` task results from the consumer.
    ///
    /// Each message is ACKed after successful deserialization.
    /// Messages that fail to deserialize are ACKed (removed) and excluded.
    pub async fn pull_task_results(&self, max_msgs: usize) -> VaporaResult<Vec<TaskResult>> {
        let batch = self.stream.pull_batch(max_msgs).await?;
        let mut results = Vec::with_capacity(batch.len());

        for (subject, payload, msg) in batch {
            match serde_json::from_slice::<TaskResult>(&payload) {
                Ok(result) => {
                    msg.ack()
                        .await
                        .map_err(|e| anyhow::anyhow!("ack failed on '{subject}': {e}"))?;
                    results.push(result);
                }
                Err(e) => {
                    tracing::warn!(subject = %subject, "dropping unparseable task result: {e}");
                    let _ = msg.ack().await;
                }
            }
        }

        Ok(results)
    }

    /// Publish an agent status heartbeat.
    pub async fn publish_agent_status(&self, status: &AgentStatus) -> VaporaResult<()> {
        let subject = format!("syntaxis.agents.{}.status", status.role);
        let payload = Bytes::from(serde_json::to_vec(status)?);
        self.stream.publish(&subject, payload).await?;
        Ok(())
    }

    /// Ping NATS by publishing a zero-byte keepalive to `syntaxis.ping`.
    pub async fn health_check(&self) -> VaporaResult<()> {
        self.stream
            .publish("syntaxis.ping", Bytes::new())
            .await
            .map_err(|e| anyhow::anyhow!("NATS health check failed: {e}"))
    }

    /// Wait up to `max_wait` for a specific `task_id` result via polling pull.
    pub async fn wait_for_completion(
        &self,
        task_id: &str,
        max_wait: Duration,
    ) -> VaporaResult<TaskResult> {
        let start = std::time::Instant::now();

        loop {
            let batch = self.pull_task_results(32).await?;
            if let Some(r) = batch.into_iter().find(|r| r.task_id == task_id) {
                return Ok(r);
            }

            if start.elapsed() > max_wait {
                return Err(anyhow::anyhow!("task '{}' timed out", task_id));
            }

            tokio::time::sleep(Duration::from_millis(200)).await;
        }
    }

    /// Return the broker configuration this bridge was created with.
    pub fn config(&self) -> &NatsBrokerConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_broker_config_default() {
        let cfg = NatsBrokerConfig::default();
        assert_eq!(cfg.url, "nats://localhost:4222");
        assert_eq!(cfg.stream_name, "SYNTAXIS");
        assert!(!cfg.require_signed_messages);
    }

    #[test]
    fn test_nats_event_config_into_broker() {
        let event_cfg = NatsEventConfig {
            url: "nats://prod:4222".to_string(),
            stream: "PROD_STREAM".to_string(),
            consumer: "prod-consumer".to_string(),
            subjects: vec!["syntaxis.>".to_string()],
            nkey_seed: Some("SXXX".to_string()),
            trusted_nkeys: vec!["PUBKEY".to_string()],
            require_signed_messages: true,
        };

        let broker: NatsBrokerConfig = event_cfg.into();
        assert_eq!(broker.url, "nats://prod:4222");
        assert_eq!(broker.stream_name, "PROD_STREAM");
        assert_eq!(broker.nkey_seed.unwrap(), "SXXX");
        assert!(broker.require_signed_messages);
    }

    #[test]
    fn test_to_platform_config_round_trip() {
        let cfg = NatsBrokerConfig {
            url: "nats://host:4222".to_string(),
            stream_name: "S".to_string(),
            consumer_name: "C".to_string(),
            subjects: vec!["s.>".to_string()],
            nkey_seed: None,
            trusted_nkeys: vec![],
            require_signed_messages: false,
        };

        let platform = cfg.to_platform_config();
        assert_eq!(platform.url, cfg.url);
        assert_eq!(platform.stream_name, cfg.stream_name);
        assert_eq!(platform.consumer_name, cfg.consumer_name);
    }

    #[test]
    fn test_task_message_serde() {
        let task = TaskMessage {
            task_id: "t1".to_string(),
            agent_role: "developer".to_string(),
            title: "Impl X".to_string(),
            description: "Details".to_string(),
            priority: 5,
            data: serde_json::json!({ "repo": "syntaxis" }),
        };
        let json = serde_json::to_string(&task).unwrap();
        let decoded: TaskMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.task_id, "t1");
        assert_eq!(decoded.priority, 5);
    }

    #[test]
    fn test_task_result_serde() {
        let r = TaskResult {
            task_id: "t1".to_string(),
            agent_role: "developer".to_string(),
            status: "completed".to_string(),
            result: serde_json::json!({ "success": true }),
            execution_time_ms: 420,
        };
        let json = serde_json::to_string(&r).unwrap();
        let decoded: TaskResult = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.status, "completed");
        assert_eq!(decoded.execution_time_ms, 420);
    }
}
