polymarket 0.1.0

use dashmap::DashSet;
use futures_util::StreamExt;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use ws_reconnect_client::{MessageStream, WebSocketClient, WebSocketClientBuilder, send_subscription, WsWriter};

use crate::{Result, websocket::{UserSubscriptionRequest, Auth, WebSocketMessage}};
use super::ChannelStream;
use super::common::spawn_forwarding_task_with_reconnect;

/// Helper function to reconnect and create a new MessageStream for the user channel
async fn reconnect_user_stream(
    client: &WebSocketClient<WebSocketMessage>,
    auth: &Auth,
    subscribed_markets: &DashSet<String>,
    writer_arc: &Arc<Mutex<Option<WsWriter>>>,
) -> Result<MessageStream<WebSocketMessage>> {
    // Establish new WebSocket connection
    let (mut writer, reader) = client.connect().await?;

    // Subscribe with current markets
    let current_markets: Vec<String> = subscribed_markets.iter().map(|e| e.clone()).collect();
    let subscription = UserSubscriptionRequest::new(auth.clone(), current_markets);
    send_subscription(&mut writer, &subscription).await?;

    // Store new writer
    *writer_arc.lock().await = Some(writer);

    // Create new MessageStream with the same shared writer Arc
    Ok(MessageStream::new(reader, writer_arc.clone(), 5))
}

/// User channel websocket client with buffered reconnection
///
/// Uses channel-based architecture for zero-message-loss reconnection:
/// - Background task polls MessageStream and sends to channel
/// - When reconnecting, spawns new task with new connection
/// - Messages buffered in channel during transition
pub struct UserWebSocket {
    client: WebSocketClient<WebSocketMessage>,
    auth: Mutex<Option<Auth>>,
    subscribed_markets: DashSet<String>,
    // Sender for spawning new stream tasks
    stream_tx: Arc<Mutex<Option<mpsc::UnboundedSender<Result<WebSocketMessage>>>>>,
    // Shared writer for sending subscription updates (same Arc used by MessageStream)
    writer: Arc<Mutex<Option<ws_reconnect_client::WsWriter>>>,
    // Task handle for the background streaming task (for cancellation on reconnect)
    task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
}

impl UserWebSocket {
    /// Create a new User WebSocket client with default configuration
    pub fn new(websocket_url: impl Into<String>) -> Self {
        let url = format!("{}/user", websocket_url.into());
        let client = WebSocketClientBuilder::new(url)
            .ping_interval(5)
            .build();

        Self {
            client,
            auth: Mutex::new(None),
            subscribed_markets: DashSet::new(),
            stream_tx: Arc::new(Mutex::new(None)),
            writer: Arc::new(Mutex::new(None)),
            task_handle: Arc::new(Mutex::new(None)),
        }
    }

    /// Create a new User WebSocket client with custom WebSocket client
    pub fn with_client(client: WebSocketClient<WebSocketMessage>) -> Self {
        Self {
            client,
            auth: Mutex::new(None),
            subscribed_markets: DashSet::new(),
            stream_tx: Arc::new(Mutex::new(None)),
            writer: Arc::new(Mutex::new(None)),
            task_handle: Arc::new(Mutex::new(None)),
        }
    }

    /// Connect to the user WebSocket and return a buffered stream
    ///
    /// This spawns a background task that forwards messages from MessageStream to a channel.
    /// The returned ChannelStream receives from this channel, providing seamless reconnection.
    ///
    /// When subscriptions change (via add_markets/remove_markets), call connect() again.
    /// The old task will naturally end, and messages are buffered in the channel during transition.
    ///
    /// # Arguments
    /// * `auth` - Authentication credentials for the user channel
    pub async fn connect(&self, auth: Auth) -> Result<ChannelStream> {
        let (tx, rx) = mpsc::unbounded_channel();

        // Establish WebSocket connection
        let (mut writer, reader) = self.client.connect().await?;

        // Store auth for reconnection
        *self.auth.lock().await = Some(auth.clone());

        // Subscribe to currently tracked markets
        let current_markets = self.get_subscribed_markets();

        let subscription = UserSubscriptionRequest::new(auth.clone(), current_markets);
        send_subscription(&mut writer, &subscription).await?;

        // Store writer in the shared Arc for both MessageStream and subscription updates
        *self.writer.lock().await = Some(writer);

        // Create MessageStream with the same shared writer Arc
        let message_stream = MessageStream::new(reader, self.writer.clone(), 5);

        // Abort any existing background task before spawning new one
        if let Some(old_handle) = self.task_handle.lock().await.take() {
            old_handle.abort();
        }

        // Spawn background task to forward messages with auto-reconnection
        let client_clone = self.client.clone();
        let auth_for_task = Arc::new(Mutex::new(Some(auth)));
        let subscribed_markets_clone = self.subscribed_markets.clone();
        let writer_arc = self.writer.clone();

        let handle = spawn_forwarding_task_with_reconnect(
            message_stream,
            tx.clone(),
            "User",
            move || {
                let client = client_clone.clone();
                let auth_mutex = auth_for_task.clone();
                let markets = subscribed_markets_clone.clone();
                let writer = writer_arc.clone();
                async move {
                    // Get auth for reconnection
                    let auth = match auth_mutex.lock().await.as_ref() {
                        Some(a) => a.clone(),
                        None => {
                            return Err(crate::PolymarketError::InvalidConfig(
                                "Cannot reconnect: no auth".to_string()
                            ).into());
                        }
                    };
                    reconnect_user_stream(&client, &auth, &markets, &writer).await
                }
            },
        );

        // Store task handle and sender for reconnection
        *self.task_handle.lock().await = Some(handle);
        *self.stream_tx.lock().await = Some(tx);

        Ok(ChannelStream { rx })
    }

    /// Add a single market ID to the subscription
    ///
    /// Convenience method for adding one market. For adding multiple markets at once,
    /// use `add_markets()` which is more efficient.
    ///
    /// # Arguments
    /// * `market_id` - Market ID to add to the subscription
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::websocket::channels::UserWebSocket;
    /// # use polymarket_sdk::websocket::Auth;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let (api_key, secret, passphrase) = ("key".to_string(), "secret".to_string(), "pass".to_string());
    /// let client = UserWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let auth = Auth::new(api_key, secret, passphrase);
    /// let mut stream = client.connect(auth).await?;
    ///
    /// // Add a single market dynamically
    /// client.add_market("market3".to_string()).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn add_market(&self, market_id: String) -> Result<()> {
        self.add_markets(vec![market_id]).await
    }

    /// Add market IDs to the subscription
    ///
    /// **IMPORTANT**: To apply changes, drop the old stream and call `connect()` again.
    /// The channel buffers messages during reconnection (< 100ms gap).
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::websocket::channels::UserWebSocket;
    /// # use polymarket_sdk::websocket::Auth;
    /// # use futures_util::StreamExt;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let (api_key, secret, passphrase) = ("key".to_string(), "secret".to_string(), "pass".to_string());
    /// let client = UserWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let auth = Auth::new(api_key, secret, passphrase);
    /// let mut stream = client.connect(auth.clone()).await?;
    ///
    /// // Add more markets
    /// client.add_markets(vec!["market3".to_string()]).await?;
    ///
    /// // Reconnect (drops old stream automatically)
    /// stream = client.connect(auth).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn add_markets(&self, market_ids: Vec<String>) -> Result<()> {
        // Add to tracking
        for market_id in market_ids {
            self.subscribed_markets.insert(market_id);
        }

        // Reconnect with updated subscriptions if already connected
        // Polymarket's WebSocket API ignores subscription updates, so we must reconnect
        if self.stream_tx.lock().await.is_some() {
            self.reconnect().await?;
        }

        Ok(())
    }

    /// Remove a single market ID from the subscription
    ///
    /// Convenience method for removing one market. For removing multiple markets at once,
    /// use `remove_markets()` which is more efficient.
    ///
    /// # Arguments
    /// * `market_id` - Market ID to remove from the subscription
    pub async fn remove_market(&self, market_id: String) -> Result<()> {
        self.remove_markets(vec![market_id]).await
    }

    /// Remove market IDs from the subscription
    ///
    /// Automatically sends a subscription update to the server if connected.
    pub async fn remove_markets(&self, market_ids: Vec<String>) -> Result<()> {
        // Remove from tracking
        for market_id in &market_ids {
            self.subscribed_markets.remove(market_id);
        }

        // Reconnect with updated subscriptions if already connected
        if self.stream_tx.lock().await.is_some() {
            self.reconnect().await?;
        }

        Ok(())
    }

    /// Reconnect with updated subscriptions while keeping the same channel
    ///
    /// This creates a new WebSocket connection with current subscriptions,
    /// spawns a new background task, and uses the existing channel to ensure
    /// the consumer's stream continues working seamlessly.
    async fn reconnect(&self) -> Result<()> {
        // Get auth (required for user channel)
        let auth = self.auth.lock().await.as_ref()
            .ok_or_else(|| crate::PolymarketError::InvalidConfig("Cannot reconnect: no auth".to_string()))?
            .clone();

        // Get the existing channel sender (must exist if we're reconnecting)
        let tx = self.stream_tx.lock().await.as_ref()
            .ok_or_else(|| crate::PolymarketError::InvalidConfig("Cannot reconnect: not connected".to_string()))?
            .clone();

        // Establish new WebSocket connection
        let (mut writer, reader) = self.client.connect().await?;

        // Subscribe with current markets
        let current_markets = self.get_subscribed_markets();
        let subscription = UserSubscriptionRequest::new(auth, current_markets);
        send_subscription(&mut writer, &subscription).await?;

        // Store new writer
        *self.writer.lock().await = Some(writer);

        // Create new MessageStream
        let mut message_stream = MessageStream::new(reader, self.writer.clone(), 5);

        // Abort old background task
        if let Some(old_handle) = self.task_handle.lock().await.take() {
            old_handle.abort();
        }

        // Spawn new background task (reusing same channel)
        let tx_clone = tx.clone();
        let handle = tokio::spawn(async move {
            while let Some(msg_result) = message_stream.next().await {
                let sdk_result = msg_result.map_err(|e| e.into());
                if tx_clone.send(sdk_result).is_err() {
                    break;
                }
            }
        });

        // Store new task handle
        *self.task_handle.lock().await = Some(handle);

        Ok(())
    }

    /// Get the currently subscribed market IDs
    ///
    /// Returns a snapshot of all market IDs that are currently subscribed.
    pub fn get_subscribed_markets(&self) -> Vec<String> {
        self.subscribed_markets
            .iter()
            .map(|entry| entry.clone())
            .collect()
    }

    /// Start listening to user messages with automatic reconnection (legacy API)
    ///
    /// This is a convenience method that combines connect() with a message handler callback.
    /// For more control, use `connect()` directly and process the stream yourself.
    ///
    /// # Arguments
    /// * `auth` - Authentication credentials for the user channel
    /// * `markets` - List of market IDs to subscribe to
    /// * `handler` - Callback function that processes each received message
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::websocket::channels::UserWebSocket;
    /// # use polymarket_sdk::websocket::Auth;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let (api_key, secret, passphrase) = ("key".to_string(), "secret".to_string(), "pass".to_string());
    /// let client = UserWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let auth = Auth::new(api_key, secret, passphrase);
    /// let markets = vec!["market_id_1".to_string()];
    ///
    /// client.listen(auth, markets, |msg| {
    ///     println!("Received user message: {:?}", msg);
    ///     Ok(())
    /// }).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn listen<F>(&self, auth: Auth, markets: Vec<String>, mut handler: F) -> Result<()>
    where
        F: FnMut(WebSocketMessage) -> Result<()>,
    {
        // Add markets first
        for market_id in markets {
            self.subscribed_markets.insert(market_id);
        }

        let subscription = UserSubscriptionRequest::new(auth, self.get_subscribed_markets());

        self.client
            .listen(Some(subscription), |msg| {
                handler(msg).map_err(|e| ws_reconnect_client::WebSocketError::HandlerError(e.to_string()))
            })
            .await?;

        Ok(())
    }
}