use anyhow::{Context, Result};
use async_nats::jetstream::{self, stream};
use bytes::Bytes;
use std::time::Duration;

// -- JetStream stream name constants --

/// Stream for wipe lifecycle events (scheduled, started, completed, failed, rolled_back).
pub const STREAM_WIPE_EVENTS: &str = "CORROSION_WIPE";

/// Stream for server telemetry (stats, status changes, crash events).
pub const STREAM_SERVER_TELEMETRY: &str = "CORROSION_TELEMETRY";

/// Stream for agent commands and responses (start, stop, file ops).
pub const STREAM_AGENT_COMMANDS: &str = "CORROSION_AGENT";

/// Stream for Steam update detection events.
pub const STREAM_STEAM_UPDATES: &str = "CORROSION_STEAM";

/// Stream for notification delivery (Discord, Pushbullet, email).
pub const STREAM_NOTIFICATIONS: &str = "CORROSION_NOTIFY";

/// Stream for license lifecycle events (activated, expired, check-in).
pub const STREAM_LICENSE_EVENTS: &str = "CORROSION_LICENSE";

/// NATS JetStream pub/sub bridge.
///
/// Centralizes all NATS interactions: stream/consumer setup, publishing
/// events, and subscribing to subjects. All services go through the
/// bridge rather than holding their own NATS client references, ensuring
/// consistent subject naming and stream configuration.
pub struct NatsBridge {
    pub client: async_nats::Client,
    pub jetstream: jetstream::Context,
}

impl NatsBridge {
    pub fn new(client: async_nats::Client) -> Self {
        let jetstream = jetstream::new(client.clone());
        Self { client, jetstream }
    }

    /// Publish a message to a NATS subject (core NATS, not JetStream).
    pub async fn publish(&self, subject: &str, payload: &[u8]) -> Result<()> {
        self.client
            .publish(subject.to_string(), Bytes::from(payload.to_vec()))
            .await
            .context("Failed to publish NATS message")?;
        Ok(())
    }

    /// Publish a message to a JetStream subject with acknowledgement.
    pub async fn publish_jetstream(&self, subject: &str, payload: &[u8]) -> Result<()> {
        self.jetstream
            .publish(subject.to_string(), Bytes::from(payload.to_vec()))
            .await
            .context("Failed to publish to JetStream")?
            .await
            .context("JetStream publish not acknowledged")?;
        Ok(())
    }

    /// Publish a JSON-serializable payload to a subject.
    pub async fn publish_json<T: serde::Serialize>(
        &self,
        subject: &str,
        payload: &T,
    ) -> Result<()> {
        let bytes = serde_json::to_vec(payload).context("Failed to serialize payload")?;
        self.publish(subject, &bytes).await
    }

    /// Publish a JSON-serializable payload to a JetStream subject.
    pub async fn publish_json_jetstream<T: serde::Serialize>(
        &self,
        subject: &str,
        payload: &T,
    ) -> Result<()> {
        let bytes = serde_json::to_vec(payload).context("Failed to serialize payload")?;
        self.publish_jetstream(subject, &bytes).await
    }

    /// Send a NATS request and wait for a reply (request/reply pattern).
    /// Used for synchronous operations like companion agent commands.
    pub async fn request(
        &self,
        subject: &str,
        payload: &[u8],
        timeout: Duration,
    ) -> Result<async_nats::Message> {
        let msg = tokio::time::timeout(
            timeout,
            self.client
                .request(subject.to_string(), Bytes::from(payload.to_vec())),
        )
        .await
        .context("NATS request timed out")?
        .context("NATS request failed")?;
        Ok(msg)
    }

    /// Send a JSON request and deserialize the reply.
    pub async fn request_json<T: serde::Serialize, R: serde::de::DeserializeOwned>(
        &self,
        subject: &str,
        payload: &T,
        timeout: Duration,
    ) -> Result<R> {
        let bytes = serde_json::to_vec(payload).context("Failed to serialize request")?;
        let reply = self.request(subject, &bytes, timeout).await?;
        serde_json::from_slice(&reply.payload).context("Failed to deserialize NATS reply")
    }

    /// Subscribe to a NATS subject and return a message stream.
    pub async fn subscribe(&self, subject: &str) -> Result<async_nats::Subscriber> {
        self.client
            .subscribe(subject.to_string())
            .await
            .context("Failed to subscribe to NATS subject")
    }

    /// Initialize all JetStream streams and consumers.
    ///
    /// Called once at application startup. Creates streams if they don't
    /// exist, updates configuration if streams already exist.
    pub async fn setup_streams(&self) -> Result<()> {
        // Wipe events: lifecycle tracking with 7-day retention
        self.ensure_stream(
            STREAM_WIPE_EVENTS,
            &["corrosion.*.wipe.>"],
            stream::RetentionPolicy::Limits,
            Duration::from_secs(7 * 24 * 3600),
        )
        .await?;

        // Server telemetry: stats/status with 24-hour retention
        self.ensure_stream(
            STREAM_SERVER_TELEMETRY,
            &["corrosion.*.telemetry.>", "corrosion.*.heartbeat", "corrosion.*.companion.heartbeat", "corrosion.*.stats"],
            stream::RetentionPolicy::Limits,
            Duration::from_secs(24 * 3600),
        )
        .await?;

        // Agent commands: reliable delivery, work queue, 1-hour TTL
        self.ensure_stream(
            STREAM_AGENT_COMMANDS,
            &["corrosion.*.cmd.>", "corrosion.*.agent.>"],
            stream::RetentionPolicy::WorkQueue,
            Duration::from_secs(3600),
        )
        .await?;

        // Steam update events: long retention for history
        self.ensure_stream(
            STREAM_STEAM_UPDATES,
            &["corrosion.steam.>"],
            stream::RetentionPolicy::Limits,
            Duration::from_secs(30 * 24 * 3600),
        )
        .await?;

        // Notification delivery queue
        self.ensure_stream(
            STREAM_NOTIFICATIONS,
            &["corrosion.*.notify.>"],
            stream::RetentionPolicy::WorkQueue,
            Duration::from_secs(24 * 3600),
        )
        .await?;

        // License lifecycle events
        self.ensure_stream(
            STREAM_LICENSE_EVENTS,
            &["corrosion.license.>"],
            stream::RetentionPolicy::Limits,
            Duration::from_secs(90 * 24 * 3600),
        )
        .await?;

        tracing::info!("JetStream streams initialized");
        Ok(())
    }

    /// Create or update a JetStream stream.
    async fn ensure_stream(
        &self,
        name: &str,
        subjects: &[&str],
        retention: stream::RetentionPolicy,
        max_age: Duration,
    ) -> Result<()> {
        let config = stream::Config {
            name: name.to_string(),
            subjects: subjects.iter().map(|s| s.to_string()).collect(),
            retention,
            max_age,
            storage: stream::StorageType::File,
            num_replicas: 1,
            ..Default::default()
        };

        self.jetstream
            .get_or_create_stream(config)
            .await
            .with_context(|| format!("Failed to create/update stream: {name}"))?;

        tracing::debug!("Stream ready: {name}");
        Ok(())
    }
}
