sockudo-adapter 4.3.1

use super::ConnectionHandler;
use bytes::Bytes;
use sockudo_core::error::Result;
use sockudo_core::websocket::SocketId;
use sockudo_protocol::ProtocolVersion;
use sockudo_protocol::constants::PONG_TIMEOUT;
use sockudo_protocol::messages::PusherMessage;
use sockudo_ws::Message;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use tracing::{debug, warn};

impl ConnectionHandler {
    pub async fn setup_initial_timeouts(
        &self,
        socket_id: &SocketId,
        app_config: &sockudo_core::app::App,
    ) -> Result<()> {
        // Set activity timeout
        self.set_activity_timeout(&app_config.id, socket_id).await?;

        // Set user authentication timeout if required
        if app_config.user_authentication_enabled() {
            let auth_timeout = self.server_options.user_authentication_timeout;
            self.set_user_authentication_timeout(&app_config.id, socket_id, auth_timeout)
                .await?;
        }

        Ok(())
    }

    pub async fn set_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
        // Clear any existing timeout
        self.clear_activity_timeout(app_id, socket_id).await?;

        let Some(connection) = self
            .connection_manager
            .get_connection(socket_id, app_id)
            .await
        else {
            return Ok(());
        };

        if connection.protocol_version == ProtocolVersion::V2 {
            debug!(
                "Skipping app-level activity timeout loop for V2 socket {} (native WebSocket ping/pong handles heartbeats)",
                socket_id
            );
            return Ok(());
        }

        let socket_id_clone = *socket_id;
        let app_id_clone = app_id.to_string();
        let connection_manager = self.connection_manager.clone();
        let handler = self.clone();
        let activity_timeout = self.server_options.activity_timeout;

        let timeout_handle = tokio::spawn(async move {
            // Initial sleep before first check
            sleep(Duration::from_secs(activity_timeout)).await;

            loop {
                // Check if connection still exists and get actual inactivity time
                let conn_manager = Arc::clone(&connection_manager);
                let conn = match conn_manager
                    .get_connection(&socket_id_clone, &app_id_clone)
                    .await
                {
                    Some(c) => c,
                    None => {
                        // Connection already cleaned up, nothing to do
                        return;
                    }
                };

                // Check actual time since last activity
                let time_since_activity = {
                    let ws = conn.inner.lock().await;
                    ws.state.time_since_last_ping()
                };

                // If less than activity timeout seconds have passed since last activity, wait more
                if time_since_activity < Duration::from_secs(activity_timeout) {
                    let remaining = Duration::from_secs(activity_timeout) - time_since_activity;
                    debug!(
                        "Socket {} still active ({}s ago), waiting {} more seconds",
                        socket_id_clone,
                        time_since_activity.as_secs(),
                        remaining.as_secs()
                    );
                    drop(conn_manager);
                    sleep(remaining).await;
                    // Continue to check again without additional delay
                    continue;
                }

                // Truly inactive for activity timeout duration, send ping
                let ping_result = {
                    let mut ws = conn.inner.lock().await;
                    if !ws.is_connected() {
                        debug!("Connection {} already closed, cleaning up", socket_id_clone);
                        drop(ws);
                        if let Err(e) = handler
                            .handle_disconnect(&app_id_clone, &socket_id_clone)
                            .await
                        {
                            warn!(
                                "Failed to handle disconnect for already closed socket {}: {}",
                                socket_id_clone, e
                            );
                        }
                        break;
                    }
                    ws.state.status = sockudo_core::websocket::ConnectionStatus::PingSent(
                        std::time::Instant::now(),
                    );
                    let ping_message = PusherMessage::ping();
                    ws.send_message(&ping_message)
                };

                match ping_result {
                    Ok(_) => {
                        debug!(
                            "Sent ping to socket {} due to activity timeout",
                            socket_id_clone
                        );

                        // Release locks before waiting for pong
                        drop(conn_manager);

                        // Wait for pong response
                        sleep(Duration::from_secs(PONG_TIMEOUT)).await;

                        // Re-acquire lock to check pong status
                        let conn_manager = Arc::clone(&connection_manager);
                        if let Some(conn) = conn_manager
                            .get_connection(&socket_id_clone, &app_id_clone)
                            .await
                        {
                            let mut ws = conn.inner.lock().await;
                            // Check if we're still in PingSent state (no pong received)
                            if let sockudo_core::websocket::ConnectionStatus::PingSent(ping_time) =
                                ws.state.status
                                && ping_time.elapsed() > Duration::from_secs(PONG_TIMEOUT)
                            {
                                warn!(
                                    "No pong received from socket {} after ping, forcing cleanup",
                                    socket_id_clone
                                );
                                let _ = ws
                                    .close(4201, "Pong reply not received in time".to_string())
                                    .await;
                                drop(ws);
                                if let Err(e) = handler
                                    .handle_disconnect(&app_id_clone, &socket_id_clone)
                                    .await
                                {
                                    warn!(
                                        "Failed to handle timeout disconnect for socket {}: {}",
                                        socket_id_clone, e
                                    );
                                }
                                break;
                            }
                        }
                        // After handling ping/pong, wait full activity timeout before next check
                        drop(conn_manager);
                        sleep(Duration::from_secs(activity_timeout)).await;
                    }
                    Err(e) => {
                        // Connection is broken (e.g., broken pipe)
                        // This is expected when client disconnects abruptly
                        debug!(
                            "Failed to send ping to socket {} (connection likely closed by client): {}",
                            socket_id_clone, e
                        );

                        if let Err(e) = handler
                            .handle_disconnect(&app_id_clone, &socket_id_clone)
                            .await
                        {
                            warn!(
                                "Failed to handle disconnect after ping failure for socket {}: {}",
                                socket_id_clone, e
                            );
                        }
                        break; // Exit the loop after cleanup
                    }
                }
            }
        });

        // Store the timeout handle
        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
            let mut ws = conn.inner.lock().await;
            ws.state.timeouts.activity_timeout_handle = Some(timeout_handle);
        }

        Ok(())
    }

    pub async fn clear_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
            let mut ws = conn.inner.lock().await;
            ws.state.timeouts.clear_activity_timeout();
        }
        Ok(())
    }

    pub async fn update_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
        // Update last activity time
        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
            let mut ws = conn.inner.lock().await;
            ws.update_activity();
        }
        Ok(())
    }

    pub async fn set_user_authentication_timeout(
        &self,
        app_id: &str,
        socket_id: &SocketId,
        timeout_seconds: u64,
    ) -> Result<()> {
        let socket_id_clone = *socket_id;
        let app_id_clone = app_id.to_string();
        let connection_manager = self.connection_manager.clone();

        // Clear any existing auth timeout
        self.clear_user_authentication_timeout(app_id, socket_id)
            .await?;

        let timeout_handle = tokio::spawn(async move {
            sleep(Duration::from_secs(timeout_seconds)).await;

            let conn_manager = connection_manager;
            if let Some(conn) = conn_manager
                .get_connection(&socket_id_clone, &app_id_clone)
                .await
            {
                let mut ws = conn.inner.lock().await;

                // Check if user is still not authenticated
                if !ws.state.is_authenticated() {
                    let _ = ws
                        .close(
                            4009,
                            "Connection not authorized within timeout.".to_string(),
                        )
                        .await;
                }
            }
        });

        // Store the timeout handle
        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
            let mut ws = conn.inner.lock().await;
            ws.state.timeouts.auth_timeout_handle = Some(timeout_handle);
        }

        Ok(())
    }

    pub async fn clear_user_authentication_timeout(
        &self,
        app_id: &str,
        socket_id: &SocketId,
    ) -> Result<()> {
        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
            let mut ws = conn.inner.lock().await;
            ws.state.timeouts.clear_auth_timeout();
        }
        Ok(())
    }

    pub async fn handle_ping_frame(
        &self,
        socket_id: &SocketId,
        app_config: &sockudo_core::app::App,
        payload: Bytes,
    ) -> Result<()> {
        // Update activity and send pong
        self.update_activity_timeout(&app_config.id, socket_id)
            .await?;

        let conn_manager = &self.connection_manager;
        if let Some(conn) = conn_manager.get_connection(socket_id, &app_config.id).await {
            let mut ws = conn.inner.lock().await;
            // Reset connection status to Active when we receive a ping (client is alive)
            ws.state.status = sockudo_core::websocket::ConnectionStatus::Active;
            // Send low-level Pong frame in response because the auto response is disabled to allow custom handling
            ws.send_frame(Message::pong(payload))?;
        }

        Ok(())
    }
}