polymarket 0.1.0

use futures_util::StreamExt;
use std::future::Future;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use ws_reconnect_client::MessageStream;

use crate::{Result, websocket::WebSocketMessage};

/// Configuration for auto-reconnection behavior
const MAX_RECONNECTION_ATTEMPTS: u32 = 10;
const BASE_BACKOFF_DELAY_MS: u64 = 1000;  // Start with 1 second
const MAX_BACKOFF_DELAY_MS: u64 = 60000;   // Cap at 60 seconds

/// Spawns a background task that forwards MessageStream messages to a channel
/// with automatic reconnection on connection loss.
///
/// This function provides the common auto-reconnection logic used by both
/// market and user WebSocket channels, eliminating code duplication.
///
/// # Type Parameters
/// * `F` - Reconnection function type (async closure)
/// * `Fut` - Future returned by reconnection function
///
/// # Arguments
/// * `message_stream` - Initial MessageStream to forward messages from
/// * `tx` - Channel sender to forward messages to
/// * `channel_name` - Name for logging (e.g., "Market" or "User")
/// * `reconnect_fn` - Async function to call when reconnection is needed
///
/// # Returns
/// A `JoinHandle` for the spawned background task
///
/// # Example
/// ```ignore
/// let handle = spawn_forwarding_task_with_reconnect(
///     message_stream,
///     tx,
///     "Market",
///     || reconnect_market_stream(&client, &tokens, &writer)
/// );
/// ```
pub(crate) fn spawn_forwarding_task_with_reconnect<F, Fut>(
    message_stream: MessageStream<WebSocketMessage>,
    tx: mpsc::UnboundedSender<Result<WebSocketMessage>>,
    channel_name: &'static str,
    reconnect_fn: F,
) -> JoinHandle<()>
where
    F: Fn() -> Fut + Send + 'static,
    Fut: Future<Output = Result<MessageStream<WebSocketMessage>>> + Send + 'static,
{
    tokio::spawn(async move {
        let mut current_stream = message_stream;
        let mut reconnect_attempt = 0u32;

        loop {
            match current_stream.next().await {
                Some(msg_result) => {
                    match msg_result {
                        Ok(msg) => {
                            // Successful message - reset reconnect counter and forward
                            reconnect_attempt = 0;

                            if tx.send(Ok(msg)).is_err() {
                                // Channel closed, stop task
                                break;
                            }
                        }
                        Err(e) => {
                            // WebSocket error - treat as connection loss and reconnect
                            eprintln!(
                                "❌ {} WebSocket error: {}, triggering reconnection",
                                channel_name, e
                            );

                            reconnect_attempt += 1;

                            if reconnect_attempt > MAX_RECONNECTION_ATTEMPTS {
                                eprintln!(
                                    "❌ {} WebSocket: Max reconnection attempts ({}) exceeded",
                                    channel_name, MAX_RECONNECTION_ATTEMPTS
                                );
                                break;
                            }

                            // Calculate exponential backoff delay: base * 2^(attempt-1), capped at max
                            let delay_ms = (BASE_BACKOFF_DELAY_MS * 2u64.pow(reconnect_attempt - 1))
                                .min(MAX_BACKOFF_DELAY_MS);

                            eprintln!(
                                "⏳ {} WebSocket: Reconnecting in {}ms (attempt {}/{})",
                                channel_name, delay_ms, reconnect_attempt, MAX_RECONNECTION_ATTEMPTS
                            );
                            tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;

                            // Attempt reconnection using provided function
                            eprintln!("🔄 {} WebSocket: Attempting to reconnect...", channel_name);
                            match reconnect_fn().await {
                                Ok(new_stream) => {
                                    current_stream = new_stream;
                                    eprintln!("✅ {} WebSocket: Reconnected successfully", channel_name);
                                }
                                Err(e) => {
                                    eprintln!(
                                        "❌ {} WebSocket: Reconnection attempt {} failed: {}",
                                        channel_name, reconnect_attempt, e
                                    );
                                    // Continue loop to retry with next backoff delay
                                }
                            }
                        }
                    }
                }
                None => {
                    // Connection died, attempt to reconnect
                    reconnect_attempt += 1;

                    if reconnect_attempt > MAX_RECONNECTION_ATTEMPTS {
                        eprintln!(
                            "❌ {} WebSocket: Max reconnection attempts ({}) exceeded",
                            channel_name, MAX_RECONNECTION_ATTEMPTS
                        );
                        break;
                    }

                    // Calculate exponential backoff delay: base * 2^(attempt-1), capped at max
                    let delay_ms = (BASE_BACKOFF_DELAY_MS * 2u64.pow(reconnect_attempt - 1))
                        .min(MAX_BACKOFF_DELAY_MS);

                    eprintln!(
                        "⏳ {} WebSocket: Reconnecting in {}ms (attempt {}/{})",
                        channel_name, delay_ms, reconnect_attempt, MAX_RECONNECTION_ATTEMPTS
                    );
                    tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;

                    // Attempt reconnection using provided function
                    eprintln!("🔄 {} WebSocket: Attempting to reconnect...", channel_name);
                    match reconnect_fn().await {
                        Ok(new_stream) => {
                            current_stream = new_stream;
                            eprintln!("✅ {} WebSocket: Reconnected successfully", channel_name);
                        }
                        Err(e) => {
                            eprintln!(
                                "❌ {} WebSocket: Reconnection attempt {} failed: {}",
                                channel_name, reconnect_attempt, e
                            );
                            // Continue loop to retry with next backoff delay
                        }
                    }
                }
            }
        }
    })
}