polymarket 0.1.0

use dashmap::DashSet;
use futures_util::{Stream, StreamExt};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
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::{MarketSubscriptionRequest, WebSocketMessage}};
use super::common::spawn_forwarding_task_with_reconnect;

/// Helper function to reconnect and create a new MessageStream for the market channel
async fn reconnect_market_stream(
    client: &WebSocketClient<WebSocketMessage>,
    subscribed_tokens: &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 tokens
    let current_tokens: Vec<String> = subscribed_tokens.iter().map(|e| e.clone()).collect();
    if !current_tokens.is_empty() {
        let subscription = MarketSubscriptionRequest::new(current_tokens);
        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))
}

/// A stream that receives messages from a channel
///
/// This wraps an unbounded receiver and implements Stream
pub struct ChannelStream {
    pub(crate) rx: mpsc::UnboundedReceiver<Result<WebSocketMessage>>,
}

impl Stream for ChannelStream {
    type Item = Result<WebSocketMessage>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.rx.poll_recv(cx)
    }
}

/// Market 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 MarketWebSocket {
    client: WebSocketClient<WebSocketMessage>,
    subscribed_tokens: 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 MarketWebSocket {
    /// Create a new Market WebSocket client with default configuration
    pub fn new(websocket_url: impl Into<String>) -> Self {
        let url = format!("{}/market", websocket_url.into());
        let client = WebSocketClientBuilder::new(url)
            .ping_interval(5)
            .build();

        Self {
            client,
            subscribed_tokens: 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 Market WebSocket client with custom WebSocket client
    pub fn with_client(client: WebSocketClient<WebSocketMessage>) -> Self {
        Self {
            client,
            subscribed_tokens: 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 market 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_tokens/remove_tokens), call connect() again.
    /// The old task will naturally end, and messages are buffered in the channel during transition.
    pub async fn connect(&self) -> Result<ChannelStream> {
        let (tx, rx) = mpsc::unbounded_channel();

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

        // Subscribe to currently tracked tokens
        let current_tokens = self.get_subscribed_tokens();

        if !current_tokens.is_empty() {
            let subscription = MarketSubscriptionRequest::new(current_tokens);
            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 subscribed_tokens_clone = self.subscribed_tokens.clone();
        let writer_arc = self.writer.clone();

        let handle = spawn_forwarding_task_with_reconnect(
            message_stream,
            tx.clone(),
            "Market",
            move || {
                let client = client_clone.clone();
                let tokens = subscribed_tokens_clone.clone();
                let writer = writer_arc.clone();
                async move { reconnect_market_stream(&client, &tokens, &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 token ID to the subscription
    ///
    /// Convenience method for adding one token. For adding multiple tokens at once,
    /// use `add_tokens()` which is more efficient.
    ///
    /// # Arguments
    /// * `token_id` - Token ID to add to the subscription
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::websocket::channels::MarketWebSocket;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = MarketWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let mut stream = client.connect().await?;
    ///
    /// // Add a single token dynamically
    /// client.add_token("token3".to_string()).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn add_token(&self, token_id: String) -> Result<()> {
        self.add_tokens(vec![token_id]).await
    }

    /// Add token 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::MarketWebSocket;
    /// # use futures_util::StreamExt;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = MarketWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let mut stream = client.connect().await?;
    ///
    /// // Add more tokens
    /// client.add_tokens(vec!["token3".to_string()]).await?;
    ///
    /// // Reconnect (drops old stream automatically)
    /// stream = client.connect().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn add_tokens(&self, token_ids: Vec<String>) -> Result<()> {
        // Add to tracking
        for token_id in token_ids {
            self.subscribed_tokens.insert(token_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(())
    }

    /// 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 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 tokens
        let current_tokens = self.get_subscribed_tokens();
        if !current_tokens.is_empty() {
            let subscription = MarketSubscriptionRequest::new(current_tokens);
            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(())
    }

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

    /// Remove token IDs from the subscription
    ///
    /// Automatically reconnects to apply the subscription changes.
    pub async fn remove_tokens(&self, token_ids: Vec<String>) -> Result<()> {
        // Remove from tracking
        for token_id in &token_ids {
            self.subscribed_tokens.remove(token_id);
        }

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

        Ok(())
    }

    /// Get the currently subscribed token IDs
    pub fn get_subscribed_tokens(&self) -> Vec<String> {
        self.subscribed_tokens
            .iter()
            .map(|entry| entry.clone())
            .collect()
    }

    /// Start listening to market 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
    /// * `token_ids` - List of token/asset IDs to subscribe to
    /// * `handler` - Callback function that processes each received message
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::websocket::channels::MarketWebSocket;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = MarketWebSocket::new("wss://ws-subscriptions-clob.polymarket.com");
    /// let token_ids = vec!["token_id_1".to_string(), "token_id_2".to_string()];
    ///
    /// client.listen(token_ids, |msg| {
    ///     println!("Received market message: {:?}", msg);
    ///     Ok(())
    /// }).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn listen<F>(&self, token_ids: Vec<String>, mut handler: F) -> Result<()>
    where
        F: FnMut(WebSocketMessage) -> Result<()>,
    {
        // Add tokens first
        for token_id in token_ids {
            self.subscribed_tokens.insert(token_id);
        }

        let subscription = MarketSubscriptionRequest::new(self.get_subscribed_tokens());

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

        Ok(())
    }
}