acta-maker-sdk 0.8.0

use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{SinkExt, StreamExt};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::net::TcpStream;

use tokio_tungstenite::{
    MaybeTlsStream, WebSocketStream, connect_async_with_config,
    tungstenite::{
        Bytes, Message, Utf8Bytes,
        protocol::{CloseFrame, WebSocketConfig},
    },
};

use crate::ws::error::{UnreadableFrame, WsClientError, WsResult, WsTransportConfigError};
use crate::ws::types::{
    AcceptQuoteMessage, AddChannelsData, AddMintsData, AuthChallengeData, BatchQuotesMessage,
    CancelAllQuotesMessage, CancelQuoteData, CancelRfqData, ClaimReferralCodeData, ClientMessage,
    CloseCode, GetActiveRfqsMessage, GetExpiriesMessage, GetIndicativePricesMessage,
    GetMakerPositionsMessage, GetMarketDescriptorsMessage, GetMarketsForMakerMessage,
    GetMarketsMessage, GetMmSummaryMessage, GetMyActiveRfqsMessage, GetMyQuotesMessage,
    GetMyReferralInfoData, GetOrderStatusMessage, GetPositionsMessage, GetSubscriptionsMessage,
    GetTokensMessage, HelloData, IndicativePricesResponseMessage, Lane, QuoteMessage,
    RedeemInviteData, RemoveChannelsData, RemoveMintsData, ReplaceQuoteMessage, ResumeAuthData,
    RfqRequestMessage, ServerMessage, StartAuthData, SubmitSignedSponsoredTxData, SubscribeData,
    UnsubscribeData, WsChannel, parse_server_message,
};
use uuid::Uuid;

type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
pub(crate) type WsSink = SplitSink<Socket, Message>;

macro_rules! ws_method {
    ($name:ident, $variant:ident, $data:ty) => {
        pub async fn $name(&mut self, data: $data) -> WsResult<()> {
            self.send(&ClientMessage::$variant(data)).await
        }
    };
}

macro_rules! ws_method_request_id {
    ($name:ident, $variant:ident, $msg:ident) => {
        pub async fn $name(&mut self) -> WsResult<()> {
            self.send(&ClientMessage::$variant($msg {
                request_id: Uuid::new_v4(),
                limit: 200,
                cursor: crate::ws::types::Cursor::Start {},
            }))
            .await
        }
    };
}

/// Low-level WebSocket transport.
///
/// [`Self::send_hello`] deliberately retains no handshake state. Callers using
/// this raw client must wait for `Welcome` and validate it with
/// Check the Welcome protocol version before sending authentication or
/// exposing the frame to application code. [`crate::ws::managed`] and the
/// vault fleet enforce that sequence themselves.
pub struct WsClient {
    stream: Socket,
    server_close: Arc<OnceLock<ServerClose>>,
    stream_label: &'static str,
}

/// Read half of a split [`WsClient`]. Liveness frames are returned, not
/// answered, so the owner decides which lane the reply takes.
pub(crate) struct WsReader {
    stream: SplitStream<Socket>,
    server_close: Arc<OnceLock<ServerClose>>,
    stream_label: &'static str,
}

/// The close frame the server sent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServerClose {
    pub code: u16,
    pub reason: String,
}

impl ServerClose {
    /// `None` for codes outside [`CloseCode`].
    #[must_use]
    pub const fn close_code(&self) -> Option<CloseCode> {
        CloseCode::from_code(self.code)
    }
}

fn record_close(slot: &OnceLock<ServerClose>, frame: Option<CloseFrame>) {
    if let Some(frame) = frame {
        let _ = slot.set(ServerClose {
            code: u16::from(frame.code),
            reason: frame.reason.to_string(),
        });
    }
}

#[allow(clippy::large_enum_variant)]
pub(crate) enum WsFrame {
    Message(ServerMessage),
    Ping(Bytes),
    Pong,
}

#[derive(Debug, Clone)]
pub struct PreparedClientMessage {
    text: Utf8Bytes,
    lane: Lane,
    subscription: Option<Box<crate::ws::types::SubscriptionChange>>,
}

impl PreparedClientMessage {
    pub fn new(message: &ClientMessage) -> Result<Self, serde_json::Error> {
        Ok(Self {
            text: serde_json::to_string(message)?.into(),
            lane: message.lane(),
            subscription: crate::ws::types::SubscriptionChange::from_message(message).map(Box::new),
        })
    }

    pub(crate) fn take_subscription(
        &mut self,
    ) -> Option<Box<crate::ws::types::SubscriptionChange>> {
        self.subscription.take()
    }

    pub(crate) const fn lane(&self) -> Lane {
        self.lane
    }

    pub(crate) fn into_text(self) -> Utf8Bytes {
        self.text
    }

    pub fn as_str(&self) -> &str {
        &self.text
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.text.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.text.is_empty()
    }
}

#[derive(Debug, Clone, Copy)]
pub struct WsTransportConfig {
    pub connect_timeout: Duration,
    pub read_buffer_size: usize,
    pub write_buffer_size: usize,
    pub max_write_buffer_size: usize,
    pub max_message_size: usize,
    pub max_frame_size: usize,
    pub tcp_nodelay: bool,
}

impl Default for WsTransportConfig {
    fn default() -> Self {
        Self {
            connect_timeout: Duration::from_secs(10),
            read_buffer_size: 64 * 1024,
            write_buffer_size: 0,
            max_write_buffer_size: 1024 * 1024,
            max_message_size: 2 * 1024 * 1024,
            max_frame_size: 2 * 1024 * 1024,
            tcp_nodelay: true,
        }
    }
}

impl WsTransportConfig {
    /// Validate limits before a connection task is spawned.
    ///
    /// # Errors
    /// Returns a typed error for zero deadlines/limits or an invalid write-buffer range.
    pub fn validate(self) -> Result<(), WsTransportConfigError> {
        if self.connect_timeout.is_zero() {
            return Err(WsTransportConfigError::ZeroConnectTimeout);
        }
        if self.max_write_buffer_size <= self.write_buffer_size {
            return Err(WsTransportConfigError::InvalidWriteBufferRange);
        }
        if self.max_message_size == 0 || self.max_frame_size == 0 {
            return Err(WsTransportConfigError::ZeroMessageOrFrameLimit);
        }
        Ok(())
    }
}

impl WsClient {
    pub async fn connect(url: &str) -> WsResult<Self> {
        Self::connect_with_config(url, WsTransportConfig::default()).await
    }

    pub async fn connect_with_config(url: &str, config: WsTransportConfig) -> WsResult<Self> {
        config.validate()?;
        let websocket_config = WebSocketConfig::default()
            .read_buffer_size(config.read_buffer_size)
            .write_buffer_size(config.write_buffer_size)
            .max_write_buffer_size(config.max_write_buffer_size)
            .max_message_size(Some(config.max_message_size))
            .max_frame_size(Some(config.max_frame_size));
        let connect = connect_async_with_config(url, Some(websocket_config), config.tcp_nodelay);
        let (stream, _response) = tokio::time::timeout(config.connect_timeout, connect)
            .await
            .map_err(|_| WsClientError::ConnectTimeout)??;
        Ok(Self {
            stream,
            server_close: Arc::default(),
            stream_label: if url.ends_with("/data") {
                "data"
            } else {
                "quote"
            },
        })
    }

    pub async fn send(&mut self, msg: &ClientMessage) -> WsResult<()> {
        let prepared = PreparedClientMessage::new(msg)?;
        self.send_prepared(&prepared).await?;
        Ok(())
    }

    pub async fn send_prepared(&mut self, message: &PreparedClientMessage) -> WsResult<()> {
        self.stream
            .send(Message::Text(message.text.clone()))
            .await?;
        Ok(())
    }

    /// Send a prepared message by transferring its buffer into the socket sink.
    pub async fn send_prepared_owned(&mut self, message: PreparedClientMessage) -> WsResult<()> {
        self.stream.send(Message::Text(message.text)).await?;
        Ok(())
    }

    ws_method!(send_hello, Hello, HelloData);
    ws_method!(resume_auth, ResumeAuth, ResumeAuthData);
    ws_method!(auth_challenge, AuthChallenge, AuthChallengeData);
    ws_method!(quote, Quote, QuoteMessage);
    ws_method!(replace_quote, ReplaceQuote, ReplaceQuoteMessage);
    ws_method!(batch_quotes, BatchQuotes, BatchQuotesMessage);
    ws_method!(subscribe, Subscribe, SubscribeData);
    ws_method!(unsubscribe, Unsubscribe, UnsubscribeData);
    ws_method!(get_my_quotes, GetMyQuotes, GetMyQuotesMessage);
    ws_method!(
        get_maker_positions,
        GetMakerPositions,
        GetMakerPositionsMessage
    );
    ws_method!(
        get_markets_for_maker,
        GetMarketsForMaker,
        GetMarketsForMakerMessage
    );
    ws_method!(rfq_request, RfqRequest, RfqRequestMessage);
    ws_method!(accept_quote, AcceptQuote, AcceptQuoteMessage);
    ws_method!(
        submit_signed_sponsored_tx,
        SubmitSignedSponsoredTx,
        SubmitSignedSponsoredTxData
    );
    ws_method!(
        indicative_prices_response,
        IndicativePricesResponse,
        IndicativePricesResponseMessage
    );
    ws_method!(
        get_indicative_prices,
        GetIndicativePrices,
        GetIndicativePricesMessage
    );
    ws_method!(
        get_market_descriptors,
        GetMarketDescriptors,
        GetMarketDescriptorsMessage
    );
    ws_method!(get_tokens, GetTokens, GetTokensMessage);
    ws_method!(get_expiries, GetExpiries, GetExpiriesMessage);
    ws_method!(get_order_status, GetOrderStatus, GetOrderStatusMessage);
    ws_method!(redeem_invite, RedeemInvite, RedeemInviteData);
    ws_method!(
        claim_referral_code,
        ClaimReferralCode,
        ClaimReferralCodeData
    );
    ws_method!(
        get_my_referral_info,
        GetMyReferralInfo,
        GetMyReferralInfoData
    );

    ws_method!(get_mm_summary, GetMmSummary, GetMmSummaryMessage);
    ws_method!(
        get_position_summary,
        GetPositionSummary,
        crate::ws::types::GetPositionSummaryMessage
    );
    ws_method_request_id!(get_subscriptions, GetSubscriptions, GetSubscriptionsMessage);
    ws_method_request_id!(get_markets, GetMarkets, GetMarketsMessage);

    ws_method!(get_positions, GetPositions, GetPositionsMessage);
    ws_method!(get_my_active_rfqs, GetMyActiveRfqs, GetMyActiveRfqsMessage);
    ws_method!(get_active_rfqs, GetActiveRfqs, GetActiveRfqsMessage);

    pub async fn ping(&mut self) -> WsResult<()> {
        self.send(&ClientMessage::Ping).await
    }

    pub async fn logout(&mut self) -> WsResult<()> {
        self.send(&ClientMessage::Logout).await
    }

    pub async fn start_auth(&mut self, pubkey: crate::Address32) -> WsResult<()> {
        self.send(&ClientMessage::StartAuth(StartAuthData { pubkey }))
            .await
    }

    pub async fn cancel_quote(&mut self, rfq_id: Uuid) -> WsResult<()> {
        self.send(&ClientMessage::CancelQuote(CancelQuoteData {
            rfq_id,
            request_id: Uuid::new_v4(),
        }))
        .await
    }

    pub async fn cancel_all_quotes(
        &mut self,
        market: crate::ws::types::QueryFilter<crate::Address32>,
    ) -> WsResult<()> {
        self.send(&ClientMessage::CancelAllQuotes(CancelAllQuotesMessage {
            request_id: Uuid::new_v4(),
            market,
        }))
        .await
    }

    pub async fn cancel_rfq(&mut self, rfq_id: Uuid) -> WsResult<()> {
        self.send(&ClientMessage::CancelRfq(CancelRfqData {
            rfq_id,
            request_id: Uuid::new_v4(),
        }))
        .await
    }

    pub async fn add_mints(
        &mut self,
        underlying_mints: Vec<crate::Address32>,
        quote_mints: Vec<crate::Address32>,
    ) -> WsResult<()> {
        self.send(&ClientMessage::AddMints(AddMintsData {
            request_id: Uuid::new_v4(),
            underlying_mints,
            quote_mints,
        }))
        .await
    }

    pub async fn remove_mints(
        &mut self,
        underlying_mints: Vec<crate::Address32>,
        quote_mints: Vec<crate::Address32>,
    ) -> WsResult<()> {
        self.send(&ClientMessage::RemoveMints(RemoveMintsData {
            request_id: Uuid::new_v4(),
            underlying_mints,
            quote_mints,
        }))
        .await
    }

    pub async fn add_channels(&mut self, channels: Vec<WsChannel>) -> WsResult<()> {
        self.send(&ClientMessage::AddChannels(AddChannelsData {
            request_id: Uuid::new_v4(),
            channels,
        }))
        .await
    }

    pub async fn remove_channels(&mut self, channels: Vec<WsChannel>) -> WsResult<()> {
        self.send(&ClientMessage::RemoveChannels(RemoveChannelsData {
            request_id: Uuid::new_v4(),
            channels,
        }))
        .await
    }

    pub async fn send_text(&mut self, text: impl Into<String>) -> WsResult<()> {
        self.stream.send(Message::Text(text.into().into())).await?;
        Ok(())
    }

    /// The server's close frame once [`Self::next`] has returned `None` after one.
    #[must_use]
    pub fn server_close(&self) -> Option<&ServerClose> {
        self.server_close.get()
    }

    pub(crate) fn server_close_slot(&self) -> Arc<OnceLock<ServerClose>> {
        Arc::clone(&self.server_close)
    }

    /// Read until the server's close frame or end of stream, at most `grace`.
    pub(crate) async fn await_server_close(&mut self, grace: Duration) {
        let _ = tokio::time::timeout(grace, async {
            while let Some(Ok(_)) = self.next().await {}
        })
        .await;
    }

    pub async fn next(&mut self) -> Option<Result<ServerMessage, WsClientError>> {
        loop {
            let msg = match self.stream.next().await {
                Some(Ok(msg)) => {
                    crate::io_observer::record(crate::io_observer::Event::Ws {
                        venue: "acta",
                        stream: self.stream_label,
                        bytes: msg.len(),
                    });
                    msg
                }
                Some(Err(err)) => return Some(Err(err.into())),
                None => return None,
            };

            match msg {
                Message::Text(text) => {
                    return Some(parse_frame(text.as_bytes()));
                }
                Message::Binary(bin) => {
                    return Some(parse_frame(&bin));
                }
                Message::Ping(payload) => {
                    if let Err(err) = self.stream.send(Message::Pong(payload)).await {
                        return Some(Err(err.into()));
                    }
                }
                Message::Pong(_) => {}
                Message::Close(frame) => {
                    record_close(&self.server_close, frame);
                    return None;
                }
                Message::Frame(_) => {}
            }
        }
    }

    pub async fn close(mut self) -> WsResult<()> {
        self.stream.close(None).await?;
        Ok(())
    }

    pub(crate) fn split(self) -> (WsSink, WsReader) {
        let (sink, stream) = self.stream.split();
        (
            sink,
            WsReader {
                stream,
                server_close: self.server_close,
                stream_label: self.stream_label,
            },
        )
    }
}

impl WsReader {
    /// Read until the server's close frame or end of stream, at most `grace`.
    pub(crate) async fn await_server_close(&mut self, grace: Duration) {
        let _ = tokio::time::timeout(grace, async {
            while let Some(Ok(_)) = self.next().await {}
        })
        .await;
    }

    pub(crate) async fn next(&mut self) -> Option<WsResult<WsFrame>> {
        loop {
            let msg = match self.stream.next().await {
                Some(Ok(msg)) => {
                    crate::io_observer::record(crate::io_observer::Event::Ws {
                        venue: "acta",
                        stream: self.stream_label,
                        bytes: msg.len(),
                    });
                    msg
                }
                Some(Err(err)) => return Some(Err(err.into())),
                None => return None,
            };
            let frame = match msg {
                Message::Text(text) => parse_frame(text.as_bytes()).map(WsFrame::Message),
                Message::Binary(bin) => parse_frame(&bin).map(WsFrame::Message),
                Message::Ping(payload) => Ok(WsFrame::Ping(payload)),
                Message::Pong(_) => Ok(WsFrame::Pong),
                Message::Close(frame) => {
                    record_close(&self.server_close, frame);
                    return None;
                }
                Message::Frame(_) => continue,
            };
            return Some(frame);
        }
    }
}

fn parse_frame(raw: &[u8]) -> WsResult<ServerMessage> {
    let parsed = match std::str::from_utf8(raw) {
        Ok(text) => parse_server_message(text),
        Err(_) => serde_json::from_slice::<ServerMessage>(raw),
    };
    parsed.map_err(|error| {
        WsClientError::UnreadableFrame(Box::new(UnreadableFrame::from_raw(raw, &error)))
    })
}