use std::time::Duration;

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

#[derive(Clone)]
pub struct NatsProducer {
    jetstream: jetstream::Context,
}

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 })
    }

    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 = if headers.is_empty() {
            self.jetstream
                .publish(subject.clone(), payload.to_vec().into())
                .await?
        } else {
            self.jetstream
                .publish_with_headers(subject, headers, payload.to_vec().into())
                .await?
        };

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

        Ok(())
    }
}

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?)
}