//! WebSocket Server for Real-time Content Updates
//!
//! Provides WebSocket connections for real-time notifications about content changes,
//! including content creation, updates, publishing status changes, and cache invalidation.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
};
use futures::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};

/// WebSocket notification types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum ContentNotification {
    /// Content was created
    ContentCreated {
        id: String,
        title: String,
        content_type: String,
        language: String,
        author: Option<String>,
    },
    /// Content was updated
    ContentUpdated {
        id: String,
        title: String,
        fields_changed: Vec<String>,
        author: Option<String>,
    },
    /// Content was published/unpublished
    ContentPublished {
        id: String,
        title: String,
        published: bool,
        author: Option<String>,
    },
    /// Content was deleted
    ContentDeleted {
        id: String,
        title: String,
        author: Option<String>,
    },
    /// Content cache was invalidated
    CacheInvalidated {
        cache_types: Vec<String>,
        affected_languages: Vec<String>,
        reason: String,
    },
    /// Content rebuild was triggered
    ContentRebuilt {
        trigger: String,
        duration_ms: u64,
        items_processed: usize,
    },
    /// System status update
    SystemStatus {
        status: String,
        message: String,
        timestamp: String,
    },
    /// Error notification
    Error {
        message: String,
        details: Option<String>,
        timestamp: String,
    },
}

/// WebSocket client information
#[derive(Debug, Clone)]
pub struct WebSocketClient {
    pub id: usize,
    pub user_agent: Option<String>,
    pub ip_address: Option<String>,
    pub connected_at: chrono::DateTime<chrono::Utc>,
    pub subscriptions: Vec<String>, // Content types or topics the client is subscribed to
}

/// WebSocket server state
pub struct WebSocketServer {
    /// Broadcast channel for sending notifications to all clients
    pub broadcast_tx: broadcast::Sender<ContentNotification>,
    /// Connected clients
    pub clients: Arc<RwLock<HashMap<usize, WebSocketClient>>>,
    /// Client ID counter
    pub client_counter: AtomicUsize,
    /// Configuration
    pub config: WebSocketConfig,
}

/// WebSocket server configuration
#[derive(Debug, Clone)]
pub struct WebSocketConfig {
    /// Maximum number of concurrent connections
    pub max_connections: usize,
    /// Send heartbeat pings every N seconds
    pub heartbeat_interval: u64,
    /// Close connection if no pong received within N seconds
    pub ping_timeout: u64,
    /// Buffer size for broadcast channel
    pub broadcast_buffer_size: usize,
}

impl Default for WebSocketConfig {
    fn default() -> Self {
        Self {
            max_connections: 100,
            heartbeat_interval: 30,
            ping_timeout: 10,
            broadcast_buffer_size: 1000,
        }
    }
}

impl WebSocketServer {
    /// Create a new WebSocket server
    pub fn new(config: WebSocketConfig) -> Self {
        let (broadcast_tx, _) = broadcast::channel(config.broadcast_buffer_size);

        Self {
            broadcast_tx,
            clients: Arc::new(RwLock::new(HashMap::new())),
            client_counter: AtomicUsize::new(0),
            config,
        }
    }

    /// Handle WebSocket upgrade request
    pub async fn handle_upgrade(
        ws: WebSocketUpgrade,
        State(server): State<Arc<WebSocketServer>>,
    ) -> Response {
        ws.on_upgrade(move |socket| server.handle_socket(socket))
    }

    /// Handle a WebSocket connection
    pub async fn handle_socket(self: Arc<Self>, socket: WebSocket) {
        let client_id = self.client_counter.fetch_add(1, Ordering::SeqCst);

        // Check connection limit
        {
            let clients = self.clients.read().await;
            if clients.len() >= self.config.max_connections {
                warn!(
                    "WebSocket connection limit reached, rejecting client {}",
                    client_id
                );
                return;
            }
        }

        info!("WebSocket client {} connected", client_id);

        // Create client record
        let client = WebSocketClient {
            id: client_id,
            user_agent: None, // Could be extracted from headers
            ip_address: None, // Could be extracted from connection
            connected_at: chrono::Utc::now(),
            subscriptions: vec!["all".to_string()], // Default to all notifications
        };

        // Add client to active connections
        {
            let mut clients = self.clients.write().await;
            clients.insert(client_id, client);
        }

        // Handle WebSocket connection in a single task
        let mut broadcast_rx = self.broadcast_tx.subscribe();
        let server_clone = self.clone();

        let connection_task = tokio::spawn(async move {
            let (mut sender, mut receiver) = socket.split();

            loop {
                tokio::select! {
                    // Handle outgoing messages (broadcast to client)
                    broadcast_result = broadcast_rx.recv() => {
                        match broadcast_result {
                            Ok(notification) => {
                                let message = match serde_json::to_string(&notification) {
                                    Ok(json) => Message::Text(json.into()),
                                    Err(e) => {
                                        error!("Failed to serialize notification: {}", e);
                                        continue;
                                    }
                                };

                                if sender.send(message).await.is_err() {
                                    debug!("Failed to send message to client {}, connection closed", client_id);
                                    break;
                                }
                            }
                            Err(_) => {
                                debug!("Broadcast channel closed for client {}", client_id);
                                break;
                            }
                        }
                    }

                    // Handle incoming messages from client
                    msg_result = receiver.next() => {
                        match msg_result {
                            Some(Ok(Message::Text(text))) => {
                                if let Err(e) = server_clone.handle_client_message(client_id, text.to_string()).await {
                                    warn!("Error handling client message from {}: {}", client_id, e);
                                }
                            }
                            Some(Ok(Message::Binary(_))) => {
                                warn!("Binary messages not supported from client {}", client_id);
                            }
                            Some(Ok(Message::Ping(ping_data))) => {
                                debug!("Received ping from client {}", client_id);
                                if sender.send(Message::Pong(ping_data)).await.is_err() {
                                    break;
                                }
                            }
                            Some(Ok(Message::Pong(_))) => {
                                debug!("Received pong from client {}", client_id);
                            }
                            Some(Ok(Message::Close(_))) => {
                                info!("Client {} requested close", client_id);
                                break;
                            }
                            Some(Err(e)) => {
                                error!("WebSocket error for client {}: {}", client_id, e);
                                break;
                            }
                            None => {
                                debug!("WebSocket stream ended for client {}", client_id);
                                break;
                            }
                        }
                    }
                }
            }
        });

        // Wait for connection task to complete
        let _ = connection_task.await;

        // Remove client from active connections
        {
            let mut clients = self.clients.write().await;
            clients.remove(&client_id);
        }

        info!("WebSocket client {} disconnected", client_id);
    }

    /// Handle incoming message from client
    async fn handle_client_message(&self, client_id: usize, message: String) -> anyhow::Result<()> {
        debug!("Received message from client {}: {}", client_id, message);

        // Parse client message
        let client_msg: ClientMessage = serde_json::from_str(&message)?;

        match client_msg {
            ClientMessage::Subscribe { topics } => {
                let mut clients = self.clients.write().await;
                if let Some(client) = clients.get_mut(&client_id) {
                    client.subscriptions = topics;
                    info!(
                        "Client {} subscribed to: {:?}",
                        client_id, client.subscriptions
                    );
                }
            }
            ClientMessage::Unsubscribe { topics } => {
                let mut clients = self.clients.write().await;
                if let Some(client) = clients.get_mut(&client_id) {
                    for topic in topics {
                        client.subscriptions.retain(|t| t != &topic);
                    }
                    info!(
                        "Client {} unsubscribed. Current subscriptions: {:?}",
                        client_id, client.subscriptions
                    );
                }
            }
            ClientMessage::Ping => {
                // Send back a system status as pong
                let notification = ContentNotification::SystemStatus {
                    status: "healthy".to_string(),
                    message: "Pong".to_string(),
                    timestamp: chrono::Utc::now().to_rfc3339(),
                };

                if let Ok(json) = serde_json::to_string(&notification) {
                    // Send directly to this client (not broadcast)
                    // For now, we'll just log it since we'd need access to the sender
                    debug!("Would send pong to client {}: {}", client_id, json);
                }
            }
        }

        Ok(())
    }

    /// Broadcast notification to all connected clients
    pub async fn broadcast_notification(&self, notification: ContentNotification) {
        debug!("Broadcasting notification: {:?}", notification);

        if let Err(e) = self.broadcast_tx.send(notification) {
            error!("Failed to broadcast notification: {}", e);
        }
    }

    /// Send notification to specific topics/subscriptions
    pub async fn send_to_subscribers(&self, notification: ContentNotification, topics: &[String]) {
        debug!(
            "Sending notification to topics {:?}: {:?}",
            topics, notification
        );

        // For now, just broadcast to all (topic filtering could be implemented later)
        if let Err(e) = self.broadcast_tx.send(notification) {
            error!("Failed to send notification to subscribers: {}", e);
        }
    }

    /// Get current connection count
    pub async fn connection_count(&self) -> usize {
        let clients = self.clients.read().await;
        clients.len()
    }

    /// Get connected clients info
    pub async fn get_clients_info(&self) -> Vec<WebSocketClient> {
        let clients = self.clients.read().await;
        clients.values().cloned().collect()
    }
}

/// Messages that clients can send to the server
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum ClientMessage {
    /// Subscribe to specific topics
    Subscribe { topics: Vec<String> },
    /// Unsubscribe from topics
    Unsubscribe { topics: Vec<String> },
    /// Ping message
    Ping,
}

/// Helper functions for creating common notifications

impl ContentNotification {
    /// Create a content created notification
    pub fn content_created(
        id: String,
        title: String,
        content_type: String,
        language: String,
        author: Option<String>,
    ) -> Self {
        Self::ContentCreated {
            id,
            title,
            content_type,
            language,
            author,
        }
    }

    /// Create a content updated notification
    pub fn content_updated(
        id: String,
        title: String,
        fields_changed: Vec<String>,
        author: Option<String>,
    ) -> Self {
        Self::ContentUpdated {
            id,
            title,
            fields_changed,
            author,
        }
    }

    /// Create a content published notification
    pub fn content_published(
        id: String,
        title: String,
        published: bool,
        author: Option<String>,
    ) -> Self {
        Self::ContentPublished {
            id,
            title,
            published,
            author,
        }
    }

    /// Create a cache invalidated notification
    pub fn cache_invalidated(
        cache_types: Vec<String>,
        affected_languages: Vec<String>,
        reason: String,
    ) -> Self {
        Self::CacheInvalidated {
            cache_types,
            affected_languages,
            reason,
        }
    }

    /// Create a system status notification
    pub fn system_status(status: String, message: String) -> Self {
        Self::SystemStatus {
            status,
            message,
            timestamp: chrono::Utc::now().to_rfc3339(),
        }
    }

    /// Create an error notification
    pub fn error(message: String, details: Option<String>) -> Self {
        Self::Error {
            message,
            details,
            timestamp: chrono::Utc::now().to_rfc3339(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_websocket_config_default() {
        let config = WebSocketConfig::default();
        assert_eq!(config.max_connections, 100);
        assert_eq!(config.heartbeat_interval, 30);
        assert_eq!(config.ping_timeout, 10);
        assert_eq!(config.broadcast_buffer_size, 1000);
    }

    #[test]
    fn test_content_notification_serialization() {
        let notification = ContentNotification::content_created(
            "test-id".to_string(),
            "Test Content".to_string(),
            "blog".to_string(),
            "en".to_string(),
            Some("author".to_string()),
        );

        let json = serde_json::to_string(&notification).unwrap();
        assert!(json.contains("ContentCreated"));
        assert!(json.contains("test-id"));
        assert!(json.contains("Test Content"));
    }

    #[tokio::test]
    async fn test_websocket_server_creation() {
        let config = WebSocketConfig::default();
        let server = WebSocketServer::new(config);

        assert_eq!(server.connection_count().await, 0);
        assert!(server.get_clients_info().await.is_empty());
    }

    #[tokio::test]
    async fn test_broadcast_notification() {
        let config = WebSocketConfig::default();
        let server = WebSocketServer::new(config);

        let notification =
            ContentNotification::system_status("test".to_string(), "Test message".to_string());

        // Should not panic even with no connected clients
        server.broadcast_notification(notification).await;
    }
}
