use std::time::Duration;

use async_nats::{HeaderMap, jetstream};
use config::AppConfig;
use serde::Serialize;
use track::CounterVec;

#[derive(Clone)]
pub struct NatsProducer {
    jetstream: jetstream::Context,
    metrics: Option<track::MetricsRegistry>,
}

impl NatsProducer {
    pub async fn new(config: &AppConfig) -> anyhow::Result<Self> {
        let jetstream = connect_jetstream(config).await?;
        ensure_stream(config, &jetstream).await?;

        Ok(Self {
            jetstream,
            metrics: None,
        })
    }

    pub fn set_metrics(&mut self, registry: track::MetricsRegistry) {
        self.metrics = Some(registry);
    }

    pub async fn send<T>(
        &self,
        subject: &str,
        key: &str,
        payload: &T,
        headers: Option<HeaderMap>,
    ) -> anyhow::Result<()>
    where
        T: Serialize + ?Sized,
    {
        let payload_bytes = serde_json::to_vec(payload)?;
        self.send_raw(subject, key, &payload_bytes, headers).await
    }

    pub async fn send_raw(
        &self,
        subject: &str,
        key: &str,
        payload: &[u8],
        headers: Option<HeaderMap>,
    ) -> anyhow::Result<()> {
        let mut headers = headers.unwrap_or_default();
        if !key.is_empty() {
            headers.append("x-message-key", key);
        }

        let subject = subject.to_string();
        let publish_result: anyhow::Result<()> = async {
            let publish = if headers.is_empty() {
                self.jetstream
                    .publish(subject.clone(), payload.to_vec().into())
                    .await?
            } else {
                self.jetstream
                    .publish_with_headers(
                        subject.clone(),
                        headers,
                        payload.to_vec().into(),
                    )
                    .await?
            };

            tokio::time::timeout(Duration::from_secs(5), publish).await??;
            Ok(())
        }
        .await;

        self.record_published(&subject, publish_result.is_ok());
        publish_result
    }

    fn record_published(&self, topic: &str, success: bool) {
        if let Some(reg) = &self.metrics {
            let status = if success { "published" } else { "error" };
            queue_messages_vec(reg)
                .with_label_values(&[topic, status])
                .inc();
        }
    }
}

pub async fn connect_jetstream(
    config: &AppConfig,
) -> anyhow::Result<jetstream::Context> {
    let client = match config.nats_token() {
        Some(token) if !token.is_empty() => {
            async_nats::ConnectOptions::with_token(token)
                .connect(config.nats_url())
                .await?
        }
        _ => async_nats::connect(config.nats_url()).await?,
    };

    Ok(jetstream::new(client))
}

pub async fn ensure_stream(
    config: &AppConfig,
    jetstream: &jetstream::Context,
) -> anyhow::Result<jetstream::stream::Stream> {
    Ok(jetstream
        .get_or_create_stream(jetstream::stream::Config {
            name: config.nats_stream_name(),
            subjects: config.nats_stream_subjects(),
            max_age: Duration::from_secs(config.nats_max_age_secs()),
            ..Default::default()
        })
        .await?)
}

fn queue_messages_vec(registry: &track::MetricsRegistry) -> CounterVec {
    registry
        .register_counter_vec(
            "queue_messages_total",
            "Total queue messages",
            &["topic", "status"],
        )
        .expect("failed to register queue_messages_total")
}