use async_nats::{
    jetstream::{self, consumer::pull, Context},
    ConnectOptions,
};
use bytes::Bytes;
use tracing::{info, warn};

use crate::{config::KeeperConfig, error::KeeperError};

/// Thin NATS + JetStream client for the keeper crates.
/// Wraps async-nats with keeper-specific subject naming conventions.
#[derive(Clone)]
pub struct KeeperNats {
    js: Context,
    workspace: String,
}

impl KeeperNats {
    pub async fn connect(cfg: &KeeperConfig) -> Result<Self, KeeperError> {
        let opts = match &cfg.nats_auth_token {
            Some(token) => ConnectOptions::new().token(token.clone()),
            None => ConnectOptions::new(),
        };

        let client = opts
            .connect(cfg.nats_url.as_str())
            .await
            .map_err(|e| KeeperError::Nats(format!("connect to {}: {e}", cfg.nats_url)))?;

        info!(url = %cfg.nats_url, workspace = %cfg.workspace, "NATS connected");
        Ok(Self {
            js: jetstream::new(client),
            workspace: cfg.workspace.clone(),
        })
    }

    /// Subscribe to ops.pending.<workspace>.> as a durable pull consumer.
    /// Uses WorkQueue semantics — only one active consumer per stream is expected.
    pub async fn pending_consumer(
        &self,
        consumer_name: &str,
    ) -> Result<pull::Stream, KeeperError> {
        let stream_name = format!("OPS_PENDING_{}", self.workspace.to_uppercase());
        let stream = self
            .js
            .get_stream(&stream_name)
            .await
            .map_err(|e| KeeperError::Nats(format!("get stream {stream_name}: {e}")))?;

        let consumer = stream
            .get_or_create_consumer(
                consumer_name,
                pull::Config {
                    durable_name: Some(consumer_name.to_string()),
                    filter_subject: format!("ops.pending.{}.>", self.workspace),
                    ack_policy: async_nats::jetstream::consumer::AckPolicy::Explicit,
                    ..Default::default()
                },
            )
            .await
            .map_err(|e| KeeperError::Nats(format!("create consumer {consumer_name}: {e}")))?;

        consumer
            .messages()
            .await
            .map_err(|e| KeeperError::Nats(format!("messages for {consumer_name}: {e}")))
    }

    /// Publish a signed JWT to ops.cmd.<workspace>.<op_type>
    pub async fn publish_signed(&self, op_type: &str, jwt: &str) -> Result<(), KeeperError> {
        let subject = format!("ops.cmd.{}.{}", self.workspace, op_type);
        self.js
            .publish(subject.clone(), Bytes::from(jwt.to_string()))
            .await
            .map_err(|e| KeeperError::Nats(format!("publish to {subject}: {e}")))?
            .await
            .map_err(|e| KeeperError::Nats(format!("ack for {subject}: {e}")))?;
        Ok(())
    }

    /// Publish an audit event JSON to ops.audit.<workspace>
    pub async fn publish_audit(&self, payload: Bytes) -> Result<(), KeeperError> {
        let subject = format!("ops.audit.{}", self.workspace);
        self.js
            .publish(subject.clone(), payload)
            .await
            .map_err(|e| KeeperError::Nats(format!("publish audit to {subject}: {e}")))?
            .await
            .map_err(|e| KeeperError::Nats(format!("audit ack for {subject}: {e}")))?;
        Ok(())
    }

    /// Fetch pending ops visible from the pending stream (used by keeper-cli list).
    /// Returns up to `max` messages without acking them.
    pub async fn peek_pending(&self, max: usize) -> Result<Vec<PeekedMessage>, KeeperError> {
        let stream_name = format!("OPS_PENDING_{}", self.workspace.to_uppercase());
        let stream = self
            .js
            .get_stream(&stream_name)
            .await
            .map_err(|e| KeeperError::Nats(format!("get stream {stream_name}: {e}")))?;

        let consumer = stream
            .get_or_create_consumer(
                "__keeper_cli_peek__",
                pull::Config {
                    durable_name: Some("__keeper_cli_peek__".to_string()),
                    ack_policy: async_nats::jetstream::consumer::AckPolicy::None,
                    ..Default::default()
                },
            )
            .await
            .map_err(|e| KeeperError::Nats(format!("peek consumer: {e}")))?;

        let mut messages_stream = consumer
            .messages()
            .await
            .map_err(|e| KeeperError::Nats(format!("peek messages: {e}")))?;

        let mut result = Vec::new();
        for _ in 0..max {
            // Non-blocking: break if no message available within a short timeout
            let msg = tokio::time::timeout(
                std::time::Duration::from_millis(200),
                futures::StreamExt::next(&mut messages_stream),
            )
            .await;
            match msg {
                Ok(Some(Ok(m))) => {
                    result.push(PeekedMessage {
                        subject: m.subject.to_string(),
                        payload: m.payload.clone(),
                    });
                }
                _ => break,
            }
        }
        Ok(result)
    }

    pub fn workspace(&self) -> &str {
        &self.workspace
    }

    /// Log a warning-level NATS state check.
    pub fn warn_if_disconnected(&self) {
        warn!("NATS connection state check requested");
    }
}

pub struct PeekedMessage {
    pub subject: String,
    pub payload: Bytes,
}
