use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::Value;
use tokio::net::TcpStream;
use tokio_tungstenite::{
connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream,
};
use crate::error::ForeciteError;
use crate::types::{FeedItem, StreamFilters};
type Ws = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamEvent {
Welcome {
tier: String,
limits: Value,
},
Feed {
#[serde(default)]
snapshot: bool,
#[serde(default)]
count: Option<u64>,
data: FeedItem,
},
QuotaExceeded {
limit: i64,
message: String,
},
Error {
code: Option<String>,
message: Option<String>,
},
#[serde(other)]
Other,
}
pub struct ForeciteStream {
ws: Ws,
}
impl ForeciteStream {
pub(crate) async fn connect(
ws_url: &str,
api_key: &str,
filters: &StreamFilters,
snapshot: u32,
) -> Result<Self, ForeciteError> {
let url = format!("{ws_url}/?api_key={}", urlencoding::encode(api_key));
let (mut ws, _) = connect_async(url).await?;
let subscribe = serde_json::json!({
"type": "subscribe",
"filters": filters,
"snapshot": snapshot,
});
ws.send(Message::Text(subscribe.to_string().into())).await?;
Ok(Self { ws })
}
pub async fn next(&mut self) -> Result<Option<StreamEvent>, ForeciteError> {
while let Some(message) = self.ws.next().await {
match message? {
Message::Text(text) => {
return Ok(Some(serde_json::from_str(&text)?));
}
Message::Ping(payload) => {
self.ws.send(Message::Pong(payload)).await?;
}
Message::Close(_) => return Ok(None),
_ => {}
}
}
Ok(None)
}
pub async fn close(mut self) -> Result<(), ForeciteError> {
self.ws.close(None).await?;
Ok(())
}
}