//! 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.  Also drives JWT lifecycle management (proactive
//! refresh, session revocation, cross-device security events) via targeted
//! per-user messages over the same channel.
//!
//! ## select! arm compatibility note
//!
//! `tokio::select!` does not support `#[cfg]` attributes on arms — they cause
//! macro parse errors.  All four arms are therefore always compiled.  In
//! non-auth builds:
//! - `refresh_timer` sleeps for `u64::MAX/2` seconds → never fires.
//! - `client_tx_keepalive` keeps the mpsc channel open but nothing ever sends →
//!   `client_rx.recv()` blocks forever → arm never fires.

use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::Duration,
};

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
};
#[cfg(feature = "auth")]
use dashmap::DashMap;
use futures::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, mpsc, RwLock};
use tracing::{debug, error, info, warn};
#[cfg(feature = "auth")]
use uuid::Uuid;

#[cfg(feature = "auth")]
use crate::auth::jwt::JwtService;

/// 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,
    },
    /// RBAC policy updated — anonymous deny-paths changed (server → WASM sync)
    RbacPolicyUpdated { anonymous_deny_paths: Vec<String> },
    /// JWT access token is approaching expiry — WASM should call
    /// server_refresh_token()
    #[cfg(feature = "auth")]
    TokenRefreshRequested,
    /// Session was explicitly revoked (logout-all or admin action)
    #[cfg(feature = "auth")]
    SessionRevoked { reason: String },
    /// A new login occurred for the same user account on another device
    #[cfg(feature = "auth")]
    NewDeviceLogin { session_id: String },
    /// A user account field was updated on another device
    #[cfg(feature = "auth")]
    AccountUpdated { field: String },
    /// Security alert — e.g. brute-force OTP attempts detected
    #[cfg(feature = "auth")]
    SecurityAlert { message: String },
    /// Push a toast notification to all connected clients.
    /// `kind` is one of "info" | "warning" | "error".
    ToastBroadcast { kind: String, message: 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>,
}

/// 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,
    /// Current RBAC deny paths for initial sync on new connections.
    pub current_deny_paths: Option<Arc<RwLock<Vec<String>>>>,
    /// Per-client targeted message senders — keyed by client_id.
    /// Allows `notify_user` to push directly to individual connections.
    #[cfg(feature = "auth")]
    pub client_senders: Arc<DashMap<usize, mpsc::Sender<ContentNotification>>>,
    /// Map from user UUID to connected client IDs for that user.
    #[cfg(feature = "auth")]
    pub user_connections: Arc<DashMap<Uuid, Vec<usize>>>,
    /// JWT service used to decode cookies at WS upgrade time.
    #[cfg(feature = "auth")]
    pub jwt_service: Option<Arc<JwtService>>,
}

/// WebSocket server configuration
#[derive(Debug, Clone)]
pub struct WebSocketConfig {
    pub max_connections: usize,
    pub heartbeat_interval: u64,
    pub ping_timeout: u64,
    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 {
    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,
            current_deny_paths: None,
            #[cfg(feature = "auth")]
            client_senders: Arc::new(DashMap::new()),
            #[cfg(feature = "auth")]
            user_connections: Arc::new(DashMap::new()),
            #[cfg(feature = "auth")]
            jwt_service: None,
        }
    }

    pub fn with_deny_paths(config: WebSocketConfig, deny_paths: Arc<RwLock<Vec<String>>>) -> Self {
        let mut server = Self::new(config);
        server.current_deny_paths = Some(deny_paths);
        server
    }

    /// Attach a JWT service so `handle_socket` can decode the session cookie
    /// at upgrade time and register user→connection mappings.
    #[cfg(feature = "auth")]
    pub fn with_jwt_service(mut self, jwt_svc: Arc<JwtService>) -> Self {
        self.jwt_service = Some(jwt_svc);
        self
    }

    /// Send a targeted notification to all active connections for `user_id`.
    #[cfg(feature = "auth")]
    pub async fn notify_user(&self, user_id: Uuid, notification: ContentNotification) {
        let client_ids: Vec<usize> = self
            .user_connections
            .get(&user_id)
            .map(|v| v.clone())
            .unwrap_or_default();

        for cid in client_ids {
            if let Some(tx) = self.client_senders.get(&cid) {
                let _ = tx.send(notification.clone()).await;
            }
        }
    }

    /// Handle WebSocket upgrade request — extracts cookies before the upgrade
    /// so they can be used inside `handle_socket` for JWT decoding.
    pub async fn handle_upgrade(
        ws: WebSocketUpgrade,
        State(server): State<Arc<WebSocketServer>>,
        headers: axum::http::HeaderMap,
    ) -> Response {
        ws.on_upgrade(move |socket| server.handle_socket(socket, headers))
    }

    /// Handle a WebSocket connection.
    ///
    /// Uses a four-arm select! loop.  Arms 3 & 4 (refresh timer and targeted
    /// notifications) are inert in non-auth builds: the timer sleeps for
    /// `u64::MAX/2` seconds and the mpsc channel has no active senders.
    pub async fn handle_socket(self: Arc<Self>, socket: WebSocket, headers: axum::http::HeaderMap) {
        // Suppress unused-variable warning in non-auth builds where headers are
        // only read by the auth-gated JWT extraction below.
        #[cfg(not(feature = "auth"))]
        let _ = &headers;

        let client_id = self.client_counter.fetch_add(1, Ordering::SeqCst);

        // Compute initial RBAC policy snapshot before the connection task.
        let initial_policy_json: Option<String> =
            if let Some(ref deny_paths_lock) = self.current_deny_paths {
                let paths = deny_paths_lock.read().await.clone();
                if !paths.is_empty() {
                    serde_json::to_string(&ContentNotification::rbac_policy_updated(paths)).ok()
                } else {
                    None
                }
            } else {
                None
            };

        {
            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);

        // ---- Always-present per-client channel ----
        // In auth builds: sender clone is registered in client_senders so notify_user()
        //   can push targeted messages; the original sender is moved into the task as
        //   a keepalive so recv() never returns None while the loop runs.
        // In non-auth builds: no registration happens; the original sender is moved
        // into   the task keeping the channel open but idle → recv() blocks
        // forever → arm 4   never fires.
        let (client_tx, mut client_rx) = mpsc::channel::<ContentNotification>(8);

        #[cfg(feature = "auth")]
        self.client_senders.insert(client_id, client_tx.clone());

        // ---- Auth: JWT decode + user→connection mapping ----
        #[cfg(feature = "auth")]
        let authenticated_user_id: Option<Uuid> =
            extract_user_id_from_upgrade_headers(&headers, &self.jwt_service);

        #[cfg(feature = "auth")]
        if let Some(uid) = authenticated_user_id {
            self.user_connections
                .entry(uid)
                .or_insert_with(Vec::new)
                .push(client_id);
        }

        // ---- Refresh timer ----
        // Auth builds: fires 120 s before JWT expiry.
        // Non-auth builds: sleeps for u64::MAX/2 seconds — effectively never fires.
        let refresh_sleep = {
            #[cfg(feature = "auth")]
            {
                compute_refresh_sleep(&self.jwt_service, &headers)
            }
            #[cfg(not(feature = "auth"))]
            {
                Duration::from_secs(u64::MAX / 2)
            }
        };
        let mut refresh_timer = Box::pin(tokio::time::sleep(refresh_sleep));

        // Register client
        {
            let mut clients = self.clients.write().await;
            clients.insert(
                client_id,
                WebSocketClient {
                    id: client_id,
                    user_agent: None,
                    ip_address: None,
                    connected_at: chrono::Utc::now(),
                    subscriptions: vec!["all".to_string()],
                },
            );
        }

        let mut broadcast_rx = self.broadcast_tx.subscribe();
        let server_clone = self.clone();

        let connection_task = tokio::spawn(async move {
            // Keep the original sender alive for the duration of the task.
            // - Auth builds: prevents the channel from closing when client_senders is
            //   cleaned up before the task ends.
            // - Non-auth builds: the only sender, so this is what keeps the channel open.
            let _tx_keepalive = client_tx;

            let (mut sender, mut receiver) = socket.split();

            if let Some(json) = initial_policy_json {
                if sender.send(Message::Text(json.into())).await.is_err() {
                    debug!("Failed to send initial policy to client {}", client_id);
                    return;
                }
            }

            loop {
                tokio::select! {
                    // Arm 1 — broadcast to all connected clients
                    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 {}, closing", client_id);
                                    break;
                                }
                            }
                            Err(_) => {
                                debug!("Broadcast channel closed for client {}", client_id);
                                break;
                            }
                        }
                    }

                    // Arm 2 — incoming client → server messages
                    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;
                            }
                        }
                    }

                    // Arm 3 — proactive JWT refresh notification
                    // Non-auth builds: refresh_timer sleeps forever, this arm never fires.
                    _ = &mut refresh_timer => {
                        #[cfg(feature = "auth")]
                        {
                            let json = serde_json::to_string(&ContentNotification::TokenRefreshRequested)
                                .unwrap_or_default();
                            let _ = sender.send(Message::Text(json.into())).await;
                        }
                        // Park the timer after firing — reconnection re-arms it.
                        refresh_timer = Box::pin(tokio::time::sleep(Duration::from_secs(u64::MAX / 2)));
                    }

                    // Arm 4 — targeted per-user notifications (session revocation, etc.)
                    // Non-auth builds: _tx_keepalive is the only sender and never sends →
                    //   this arm blocks indefinitely and never fires.
                    Some(notif) = client_rx.recv() => {
                        match serde_json::to_string(&notif) {
                            Ok(json) => {
                                if sender.send(Message::Text(json.into())).await.is_err() {
                                    debug!("Failed to send targeted notification to client {}, closing", client_id);
                                    break;
                                }
                            }
                            Err(e) => error!("Failed to serialize targeted notification: {}", e),
                        }
                    }
                }
            }
        });

        let _ = connection_task.await;

        // ---- Cleanup ----
        #[cfg(feature = "auth")]
        {
            self.client_senders.remove(&client_id);
            if let Some(uid) = authenticated_user_id {
                if let Some(mut entry) = self.user_connections.get_mut(&uid) {
                    entry.retain(|&id| id != client_id);
                }
            }
        }

        {
            let mut clients = self.clients.write().await;
            clients.remove(&client_id);
        }

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

    async fn handle_client_message(&self, client_id: usize, message: String) -> anyhow::Result<()> {
        debug!("Received message from client {}: {}", client_id, 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 => {
                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) {
                    debug!("Would send pong to client {}: {}", client_id, json);
                }
            }
            ClientMessage::RouteView { path, language } => {
                debug!(
                    "Client {} navigated to '{}' (language: '{}')",
                    client_id, path, language
                );
            }
        }

        Ok(())
    }

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

    pub async fn send_to_subscribers(&self, notification: ContentNotification, topics: &[String]) {
        debug!(
            "Sending notification to topics {:?}: {:?}",
            topics, notification
        );
        if let Err(e) = self.broadcast_tx.send(notification) {
            error!("Failed to send notification to subscribers: {}", e);
        }
    }

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

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

// ---------------------------------------------------------------------------
// Auth-gated helper functions
// ---------------------------------------------------------------------------

/// Extract the authenticated user UUID from the HTTP Upgrade request's cookies.
/// Returns `None` if no JWT service is configured, no session cookie is
/// present, or the token fails validation.
#[cfg(feature = "auth")]
fn extract_user_id_from_upgrade_headers(
    headers: &axum::http::HeaderMap,
    jwt_service: &Option<Arc<JwtService>>,
) -> Option<Uuid> {
    let svc = jwt_service.as_ref()?;
    let cookie_hdr = headers.get(axum::http::header::COOKIE)?.to_str().ok()?;

    let token = cookie_hdr
        .split(';')
        .find_map(|pair| pair.trim().strip_prefix("session_token="))?
        .to_string();

    let claims = svc.verify_access_token(&token).ok()?;
    claims.sub.parse::<Uuid>().ok()
}

/// Compute how long to sleep before sending `TokenRefreshRequested`.
/// Targets `exp − 120 s`; returns 1 s when the token is already near expiry.
/// Returns `u64::MAX / 2` when no JWT service is configured or no token cookie
/// exists.
#[cfg(feature = "auth")]
fn compute_refresh_sleep(
    jwt_service: &Option<Arc<JwtService>>,
    headers: &axum::http::HeaderMap,
) -> Duration {
    let Some(svc) = jwt_service.as_ref() else {
        return Duration::from_secs(u64::MAX / 2);
    };

    let cookie_hdr = headers
        .get(axum::http::header::COOKIE)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");

    let token = cookie_hdr
        .split(';')
        .find_map(|pair| pair.trim().strip_prefix("session_token="))
        .unwrap_or("");

    let Ok(claims) = svc.verify_access_token(token) else {
        return Duration::from_secs(u64::MAX / 2);
    };

    let now_secs = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    let exp = claims.exp as u64;
    if exp <= now_secs + 120 {
        Duration::from_secs(1)
    } else {
        Duration::from_secs(exp - now_secs - 120)
    }
}

// ---------------------------------------------------------------------------
// Client → Server messages
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum ClientMessage {
    Subscribe { topics: Vec<String> },
    Unsubscribe { topics: Vec<String> },
    Ping,
    RouteView { path: String, language: String },
}

// ---------------------------------------------------------------------------
// ContentNotification constructor helpers
// ---------------------------------------------------------------------------

impl ContentNotification {
    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,
        }
    }

    pub fn content_updated(
        id: String,
        title: String,
        fields_changed: Vec<String>,
        author: Option<String>,
    ) -> Self {
        Self::ContentUpdated {
            id,
            title,
            fields_changed,
            author,
        }
    }

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

    pub fn cache_invalidated(
        cache_types: Vec<String>,
        affected_languages: Vec<String>,
        reason: String,
    ) -> Self {
        Self::CacheInvalidated {
            cache_types,
            affected_languages,
            reason,
        }
    }

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

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

    pub fn rbac_policy_updated(anonymous_deny_paths: Vec<String>) -> Self {
        Self::RbacPolicyUpdated {
            anonymous_deny_paths,
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[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());

        server.broadcast_notification(notification).await;
    }

    #[cfg(feature = "auth")]
    #[test]
    fn test_auth_notification_serialization() {
        let notif = ContentNotification::SessionRevoked {
            reason: "logout_all".to_string(),
        };
        let json = serde_json::to_string(&notif).unwrap();
        assert!(json.contains("SessionRevoked"));
        assert!(json.contains("logout_all"));

        let notif2 = ContentNotification::SecurityAlert {
            message: "Too many OTP attempts".to_string(),
        };
        let json2 = serde_json::to_string(&notif2).unwrap();
        assert!(json2.contains("SecurityAlert"));
    }
}
