//! SSE Notifier for managing client subscriptions
//!
//! This module provides the subscription management infrastructure for SSE streaming.
//! The `SseNotifier` manages per-client broadcast channels for targeted notifications.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::{Arc, RwLock};

use tokio::sync::broadcast;
use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::Stream;
use tracing::{debug, trace, warn};

use super::types::SessionUpdateParams;

/// Default broadcast channel capacity
const DEFAULT_CHANNEL_CAPACITY: usize = 256;

/// Internal state for SSE subscriptions
#[derive(Debug, Default)]
struct SseNotifierState {
    /// Map from client_id to their broadcast sender
    subscriptions: HashMap<String, broadcast::Sender<SessionUpdateParams>>,
}

/// SSE Notifier for managing client subscriptions and broadcasting notifications
///
/// The `SseNotifier` manages per-client subscriptions using tokio broadcast channels.
/// Each client has their own broadcast channel, allowing targeted notifications.
///
/// ## Thread Safety
///
/// The `SseNotifier` is designed to be cloned and shared across async tasks.
/// Internal state is protected by `Arc<RwLock<...>>` for thread-safe access.
///
/// ## Example
///
/// ```rust,ignore
/// let notifier = SseNotifier::new();
///
/// // Subscribe a client and get their stream
/// let stream = notifier.subscribe("client-123");
///
/// // In another task, broadcast notifications
/// notifier.broadcast("client-123", notification);
///
/// // When client disconnects
/// notifier.unsubscribe("client-123");
/// ```
#[derive(Debug, Clone)]
pub struct SseNotifier {
    state: Arc<RwLock<SseNotifierState>>,
    channel_capacity: usize,
}

impl Default for SseNotifier {
    fn default() -> Self {
        Self::new()
    }
}

impl SseNotifier {
    /// Create a new SSE notifier with default capacity
    ///
    /// Creates a new notifier with the default channel capacity of 256 messages.
    pub fn new() -> Self {
        Self {
            state: Arc::new(RwLock::new(SseNotifierState::default())),
            channel_capacity: DEFAULT_CHANNEL_CAPACITY,
        }
    }

    /// Create a new SSE notifier with custom channel capacity
    ///
    /// # Parameters
    ///
    /// - `capacity`: The maximum number of messages that can be buffered per client
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            state: Arc::new(RwLock::new(SseNotifierState::default())),
            channel_capacity: capacity,
        }
    }

    /// Subscribe a client and return an async stream of notifications
    ///
    /// Creates a broadcast channel for the client and returns a stream that can
    /// be used with Axum's `Sse<impl Stream>` response type.
    ///
    /// If the client is already subscribed, a new receiver is created from the
    /// existing sender (allowing multiple connections per client).
    ///
    /// # Parameters
    ///
    /// - `client_id`: Unique identifier for the client
    ///
    /// # Returns
    ///
    /// A pinned stream of `SessionUpdateParams` wrapped in `Result` for error handling.
    /// The stream is compatible with Axum's SSE handler.
    pub fn subscribe(
        &self,
        client_id: &str,
    ) -> Pin<Box<dyn Stream<Item = Result<SessionUpdateParams, BroadcastStreamRecvError>> + Send>>
    {
        let mut state = self.state.write().expect("Lock poisoned");

        // Get or create the sender for this client
        let sender = state
            .subscriptions
            .entry(client_id.to_string())
            .or_insert_with(|| {
                debug!("Creating new broadcast channel for client: {}", client_id);
                let (tx, _rx) = broadcast::channel(self.channel_capacity);
                tx
            });

        // Create a new receiver
        let receiver = sender.subscribe();

        debug!(
            "Client subscribed to SSE: {} (subscribers: {})",
            client_id,
            sender.receiver_count()
        );

        // Convert to stream using tokio_stream
        let stream = BroadcastStream::new(receiver);

        Box::pin(stream)
    }

    /// Unsubscribe a client and cleanup resources
    ///
    /// Removes the client's subscription and drops the sender, which will
    /// cause all associated receivers to receive an error on their next poll.
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client to unsubscribe
    ///
    /// # Returns
    ///
    /// `true` if the client was subscribed and removed, `false` if not found
    pub fn unsubscribe(&self, client_id: &str) -> bool {
        let mut state = self.state.write().expect("Lock poisoned");

        if let Some(sender) = state.subscriptions.remove(client_id) {
            debug!(
                "Client unsubscribed from SSE: {} (was {} receivers)",
                client_id,
                sender.receiver_count()
            );
            // Sender is dropped here, which will close the channel
            true
        } else {
            debug!("Client was not subscribed: {}", client_id);
            false
        }
    }

    /// Broadcast a session update to a specific client
    ///
    /// Sends the session update to all receivers subscribed to the specified client.
    /// If no receivers are active (e.g., all have dropped or lagged), the send
    /// is handled gracefully.
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client to send to
    /// - `update_params`: The session update parameters to send
    ///
    /// # Returns
    ///
    /// - `Ok(n)` where n is the number of receivers that received the update
    /// - `Err(())` if the client is not subscribed
    pub fn broadcast(
        &self,
        client_id: &str,
        update_params: SessionUpdateParams,
    ) -> Result<usize, ()> {
        let state = self.state.read().expect("Lock poisoned");

        if let Some(sender) = state.subscriptions.get(client_id) {
            match sender.send(update_params) {
                Ok(n) => {
                    trace!(
                        "Broadcast session update to client {}: {} receivers",
                        client_id,
                        n
                    );
                    Ok(n)
                }
                Err(_) => {
                    // No active receivers - this is okay, the update is just dropped
                    trace!("Broadcast to client {}: no active receivers", client_id);
                    Ok(0)
                }
            }
        } else {
            warn!("Attempted broadcast to unsubscribed client: {}", client_id);
            Err(())
        }
    }

    /// Broadcast a session update to all subscribed clients
    ///
    /// Sends the session update to all connected clients. Failed sends to
    /// individual clients are logged but don't stop the broadcast.
    ///
    /// # Parameters
    ///
    /// - `update_params`: The session update parameters to broadcast
    ///
    /// # Returns
    ///
    /// The total number of receivers that received the update
    pub fn broadcast_all(&self, update_params: SessionUpdateParams) -> usize {
        let state = self.state.read().expect("Lock poisoned");
        let mut total_receivers = 0;

        for (client_id, sender) in state.subscriptions.iter() {
            match sender.send(update_params.clone()) {
                Ok(n) => {
                    total_receivers += n;
                    trace!("Broadcast to client {}: {} receivers", client_id, n);
                }
                Err(_) => {
                    trace!("Broadcast to client {}: no active receivers", client_id);
                }
            }
        }

        debug!(
            "Broadcast to all clients: {} total receivers across {} clients",
            total_receivers,
            state.subscriptions.len()
        );

        total_receivers
    }

    /// Get the number of subscribed clients
    pub fn client_count(&self) -> usize {
        let state = self.state.read().expect("Lock poisoned");
        state.subscriptions.len()
    }

    /// Check if a client is subscribed
    pub fn is_subscribed(&self, client_id: &str) -> bool {
        let state = self.state.read().expect("Lock poisoned");
        state.subscriptions.contains_key(client_id)
    }

    /// Get the list of subscribed client IDs
    pub fn subscribed_clients(&self) -> Vec<String> {
        let state = self.state.read().expect("Lock poisoned");
        state.subscriptions.keys().cloned().collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sse::types::SessionUpdateVariant;
    use tokio_stream::StreamExt;

    #[test]
    fn test_sse_notifier_new() {
        let notifier = SseNotifier::new();
        assert_eq!(notifier.client_count(), 0);
    }

    #[test]
    fn test_sse_notifier_with_capacity() {
        let notifier = SseNotifier::with_capacity(512);
        assert_eq!(notifier.channel_capacity, 512);
    }

    #[tokio::test]
    async fn test_subscribe_unsubscribe() {
        let notifier = SseNotifier::new();

        // Subscribe
        let _stream = notifier.subscribe("client-1");
        assert!(notifier.is_subscribed("client-1"));
        assert_eq!(notifier.client_count(), 1);

        // Unsubscribe
        assert!(notifier.unsubscribe("client-1"));
        assert!(!notifier.is_subscribed("client-1"));
        assert_eq!(notifier.client_count(), 0);

        // Unsubscribe non-existent
        assert!(!notifier.unsubscribe("client-1"));
    }

    #[tokio::test]
    async fn test_broadcast_to_client() {
        let notifier = SseNotifier::new();

        // Subscribe
        let mut stream = notifier.subscribe("client-1");

        // Broadcast
        let update_params = SessionUpdateParams {
            session_id: "sess-1".to_string(),
            update: SessionUpdateVariant::SessionIdle {},
            ..Default::default()
        };

        let result = notifier.broadcast("client-1", update_params.clone());
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 1);

        // Receive
        let received = stream.next().await;
        assert!(received.is_some());
        let received = received.unwrap();
        assert!(received.is_ok());
        assert_eq!(received.unwrap(), update_params);
    }

    #[tokio::test]
    async fn test_broadcast_to_unsubscribed_client() {
        let notifier = SseNotifier::new();

        let update_params = SessionUpdateParams {
            session_id: "sess-1".to_string(),
            update: SessionUpdateVariant::SessionIdle {},
            ..Default::default()
        };

        let result = notifier.broadcast("unknown-client", update_params);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_broadcast_all() {
        let notifier = SseNotifier::new();

        // Subscribe multiple clients
        let mut stream1 = notifier.subscribe("client-1");
        let mut stream2 = notifier.subscribe("client-2");

        let update_params = SessionUpdateParams {
            session_id: "sess-1".to_string(),
            update: SessionUpdateVariant::SessionIdle {},
            ..Default::default()
        };

        let count = notifier.broadcast_all(update_params.clone());
        assert_eq!(count, 2);

        // Both should receive
        let r1 = stream1.next().await.unwrap().unwrap();
        let r2 = stream2.next().await.unwrap().unwrap();

        assert_eq!(r1, update_params);
        assert_eq!(r2, update_params);
    }

    #[test]
    fn test_subscribed_clients() {
        let notifier = SseNotifier::new();

        let _s1 = notifier.subscribe("client-1");
        let _s2 = notifier.subscribe("client-2");

        let clients = notifier.subscribed_clients();
        assert_eq!(clients.len(), 2);
        assert!(clients.contains(&"client-1".to_string()));
        assert!(clients.contains(&"client-2".to_string()));
    }

    #[tokio::test]
    async fn test_multiple_receivers_same_client() {
        let notifier = SseNotifier::new();

        // Subscribe same client twice
        let mut stream1 = notifier.subscribe("client-1");
        let mut stream2 = notifier.subscribe("client-1");

        let update_params = SessionUpdateParams {
            session_id: "sess-1".to_string(),
            update: SessionUpdateVariant::SessionIdle {},
            ..Default::default()
        };

        let count = notifier.broadcast("client-1", update_params.clone());
        assert!(count.is_ok());
        assert_eq!(count.unwrap(), 2); // Both receivers

        // Both should receive
        let r1 = stream1.next().await.unwrap().unwrap();
        let r2 = stream2.next().await.unwrap().unwrap();

        assert_eq!(r1, update_params);
        assert_eq!(r2, update_params);
    }
}
