use std::time::Duration;

use async_nats::{
    jetstream::{self, consumer::pull, publish::PublishAck, Context},
    ConnectOptions,
};
use bytes::Bytes;
use tokio::time::sleep;
use tracing::{info, warn};

use crate::{config::NatsConfig, streams::ensure_provisioning_streams, Error};

/// Shared NATS JetStream bridge.
///
/// Clone-cheap — both fields are internally `Arc`-backed.
#[derive(Clone)]
pub struct NatsBridge {
    pub(crate) client: async_nats::Client,
    pub(crate) jetstream: Context,
    prefix: String,
}

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

impl NatsBridge {
    /// Connect to NATS with exponential backoff retry.
    /// Ensures all provisioning JetStream streams exist before returning.
    pub async fn connect(config: &NatsConfig) -> Result<Self, Error> {
        let url = config.url.clone();
        let auth_token = config.auth_token.clone();
        let max_reconnects = config.max_reconnects;
        let reconnect_wait_ms = config.reconnect_wait_ms;

        let mut last_err = None;
        for attempt in 0..=max_reconnects {
            let mut opts = ConnectOptions::new();
            if let Some(ref token) = auth_token {
                opts = opts.token(token.clone());
            }

            match opts.connect(url.as_str()).await {
                Ok(client) => {
                    info!(url = %url, "NATS connected");
                    let js = jetstream::new(client.clone());
                    ensure_provisioning_streams(&js).await?;
                    return Ok(Self {
                        client,
                        jetstream: js,
                        prefix: config.subject_prefix.clone(),
                    });
                }
                Err(e) => {
                    last_err = Some(e.to_string());
                    if attempt < max_reconnects {
                        let wait_ms = reconnect_wait_ms * (2u64.pow(attempt.min(6)));
                        warn!(
                            attempt = attempt + 1,
                            max = max_reconnects,
                            wait_ms,
                            "NATS connect failed, retrying"
                        );
                        sleep(Duration::from_millis(wait_ms)).await;
                    }
                }
            }
        }

        Err(Error::Connect(last_err.unwrap_or_else(|| "unknown".into())))
    }

    /// Publish a message to a subject relative to the configured prefix.
    ///
    /// `subject = "tasks.submitted"` → publishes to
    /// `"provisioning.tasks.submitted"`. Returns the `PublishAck` after
    /// JetStream confirms persistence.
    pub async fn publish(&self, subject: &str, payload: Bytes) -> Result<PublishAck, Error> {
        let full = format!("{}.{}", self.prefix, subject);
        self.jetstream
            .publish(full, payload)
            .await
            .map_err(|e| Error::Publish(e.to_string()))?
            .await
            .map_err(|e| Error::Publish(e.to_string()))
    }

    /// Publish JSON-serialisable value to a subject.
    pub async fn publish_json<T: serde::Serialize>(
        &self,
        subject: &str,
        value: &T,
    ) -> Result<PublishAck, Error> {
        let payload = serde_json::to_vec(value).map_err(|e| Error::Serialization(e.to_string()))?;
        self.publish(subject, Bytes::from(payload)).await
    }

    /// Obtain a pull consumer message stream for durable consumption.
    ///
    /// Creates the consumer if it does not yet exist (idempotent).
    pub async fn subscribe_pull(
        &self,
        stream: &str,
        consumer_name: &str,
    ) -> Result<pull::Stream, Error> {
        let js_stream = self
            .jetstream
            .get_stream(stream)
            .await
            .map_err(|e| Error::Stream(format!("get stream {stream}: {e}")))?;

        let consumer = js_stream
            .get_or_create_consumer(
                consumer_name,
                pull::Config {
                    durable_name: Some(consumer_name.to_string()),
                    ack_policy: async_nats::jetstream::consumer::AckPolicy::Explicit,
                    ..Default::default()
                },
            )
            .await
            .map_err(|e| Error::Consumer(format!("{consumer_name}: {e}")))?;

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

    /// Check connectivity (NATS connection state).
    pub fn health(&self) -> Result<(), Error> {
        match self.client.connection_state() {
            async_nats::connection::State::Connected => Ok(()),
            state => Err(Error::Health(format!("NATS state: {state:?}"))),
        }
    }

    /// Access the raw JetStream context for advanced operations.
    pub fn jetstream(&self) -> &Context {
        &self.jetstream
    }

    /// Subscribe to a core NATS subject (not JetStream).
    ///
    /// Used by services that need push-based subscriptions for internal
    /// coordination subjects not backed by JetStream persistence.
    pub async fn subscribe(&self, subject: impl Into<String>) -> Result<async_nats::Subscriber, Error> {
        self.client
            .subscribe(subject.into())
            .await
            .map_err(|e| Error::Connect(format!("subscribe failed: {e}")))
    }

    /// Publish a raw message to a core NATS subject (not JetStream).
    pub async fn publish_raw(&self, subject: &str, payload: bytes::Bytes) -> Result<(), Error> {
        self.client
            .publish(subject.to_string(), payload)
            .await
            .map_err(|e| Error::Publish(e.to_string()))
    }
}
