use std::fmt::{self, Debug};
use std::sync::Arc;
use std::vec;
use crate::error::{Error, Result};
use crate::event::*;
use crate::process::Processor;
use async_trait::async_trait;
use dashmap::DashMap;
use futures_util::stream::SplitSink;
use futures_util::{SinkExt, StreamExt};
use serde_derive::{Deserialize, Serialize};
use serde_json::json;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::time::{Duration, sleep};
use tokio_tungstenite::{WebSocketStream, accept_async, connect_async};
#[derive(Debug)]
pub struct BotContext {
connection: Mutex<Option<BotConnection>>,
pub url: Option<String>,
pub id: i64,
pub processors: Arc<Vec<Processor>>,
pub echo_notifer: Arc<DashMap<String, tokio::sync::mpsc::Sender<Response>>>,
}
pub struct EchoAsyncResponse(
String,
tokio::sync::mpsc::Receiver<Response>,
Arc<DashMap<String, tokio::sync::mpsc::Sender<Response>>>,
);
impl Drop for EchoAsyncResponse {
fn drop(&mut self) {
self.2.remove(&self.0);
}
}
impl EchoAsyncResponse {
pub async fn response(mut self, timeout: Duration) -> Result<Response> {
let r = tokio::time::timeout(timeout, async { self.1.recv().await }).await?;
Ok(r.ok_or(Error::StateError("response not received".to_string()))?)
}
pub async fn data(self, timeout: Duration) -> Result<serde_json::Value> {
let r = self.response(timeout).await?;
if r.retcode != 0 {
return Err(Error::StateError(r.message));
} else {
Ok(r.data)
}
}
}
pub struct SendMessageAsyncResponse(EchoAsyncResponse);
#[derive(Clone, Copy, Serialize, Deserialize, Debug)]
pub struct SendMessageResponse {
pub message_id: i64,
}
impl SendMessageAsyncResponse {
pub async fn wait_response_with_timeout(
self,
timeout: Duration,
) -> Result<SendMessageResponse> {
Ok(serde_json::from_value(self.0.data(timeout).await?)?)
}
pub async fn wait_response(self) -> Result<SendMessageResponse> {
self.wait_response_with_timeout(Duration::from_secs(10))
.await
}
}
impl BotContext {
pub async fn websocket_send(
&self,
action: &str,
msg: serde_json::Value,
) -> Result<EchoAsyncResponse> {
let echo = uuid::Uuid::new_v4().to_string();
let (sender, receiver) = tokio::sync::mpsc::channel::<Response>(1);
self.echo_notifer.insert(echo.clone(), sender);
let echo_response = EchoAsyncResponse(echo.clone(), receiver, self.echo_notifer.clone());
let msg = json!(
{
"action": action,
"params": msg,
"echo": echo
}
);
let msg = serde_json::to_string(&msg).unwrap();
tracing::debug!("WS send: {}", msg);
let mut connection_lock = self.connection.lock().await;
let connection = connection_lock
.as_mut()
.ok_or(Error::StateError("connection not ready".to_string()))?;
connection.send_raw(msg).await?;
Ok(echo_response)
}
pub async fn send_private_message(
&self,
user_id: i64,
message: impl SendMessage,
) -> Result<SendMessageAsyncResponse> {
let msg = json!(
{
"user_id": user_id,
"message": message.json()?,
}
);
self.websocket_send("send_private_msg", msg)
.await
.map(|r| SendMessageAsyncResponse(r))
}
pub async fn send_group_message(
&self,
group_id: i64,
message: impl SendMessage,
) -> Result<SendMessageAsyncResponse> {
let msg = json!(
{
"group_id": group_id,
"message": message.json()?,
}
);
self.websocket_send("send_group_msg", msg)
.await
.map(|r| SendMessageAsyncResponse(r))
}
pub async fn send_message(
&self,
message_type: MessageType,
target_id: i64,
message: impl SendMessage,
) -> Result<SendMessageAsyncResponse> {
match message_type {
MessageType::Private => self.send_private_message(target_id, message).await,
MessageType::Group => self.send_group_message(target_id, message).await,
_ => Err(Error::FieldError("unknown message_type".to_string())),
}
}
pub async fn delete_msg(&self, message_id: i64) -> Result<EchoAsyncResponse> {
let msg = json!(
{
"message_id": message_id,
}
);
self.websocket_send("delete_msg", msg).await
}
pub async fn get_msg_with_timeout(
&self,
message_id: i64,
timeout: Duration,
) -> Result<Message> {
let send = json!({
"message_id": message_id,
});
let response = self.websocket_send("get_msg", send).await?;
let data = response.data(timeout).await?;
let msg = Message::parse(&data)?;
Ok(msg)
}
pub async fn get_msg(&self, message_id: i64) -> Result<Message> {
self.get_msg_with_timeout(message_id, Duration::from_secs(3))
.await
}
pub async fn get_forward_msg_with_timeout(
&self,
id: &str,
timeout: Duration,
) -> Result<ForwardMessage> {
let send = json!({
"id": id,
});
let response = self.websocket_send("get_forward_msg", send).await?;
let data = response.data(timeout).await?;
let msg = ForwardMessage::parse(&data)?;
Ok(msg)
}
pub async fn get_forward_msg(&self, id: &str) -> Result<ForwardMessage> {
self.get_forward_msg_with_timeout(id, Duration::from_secs(3))
.await
}
}
impl BotContext {
async fn set_connection(&self, connection: impl Into<Option<BotConnection>>) {
let mut connection_lock = self.connection.lock().await;
*connection_lock = connection.into();
}
async fn handle_receive(
&self,
bot_ctx: Arc<BotContext>,
msg: &tokio_tungstenite::tungstenite::protocol::Message,
) {
match msg {
tokio_tungstenite::tungstenite::protocol::Message::Text(text) => {
tracing::debug!("WS received: {}", text.to_string());
match parse_post(text) {
Ok(post) => {
tracing::debug!("parse post: {:?}", post);
for processor in self.processors.iter() {
let processe_result = processor.process(bot_ctx.clone(), &post).await;
match processe_result {
Ok(b) => {
if b {
break;
}
},
Err(err) => {
tracing::error!("processor {:?} process post {:?} error: {:?}", processor, post, err);
break;
},
}
}
}
Err(e) => {
tracing::error!("WS received: {:?}", e);
}
}
}
_ => {
tracing::error!("WS received: {:?}", msg);
}
}
}
}
#[async_trait]
pub trait WsWriter {
async fn send_raw(&mut self, msg: String) -> Result<()>;
}
#[async_trait]
impl<S> WsWriter
for SplitSink<WebSocketStream<S>, tokio_tungstenite::tungstenite::protocol::Message>
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
async fn send_raw(&mut self, msg: String) -> Result<()> {
self.send(tokio_tungstenite::tungstenite::protocol::Message::Text(
msg.into(),
))
.await?;
Ok(())
}
}
struct BotConnection {
sender: Box<dyn WsWriter + Send + Sync>,
}
impl BotConnection {
pub async fn send_raw(&mut self, msg: String) -> Result<()> {
self.sender.send_raw(msg).await?;
Ok(())
}
}
impl fmt::Debug for BotConnection {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BotConnection")
.field("sender", &"Box<dyn WsWriter>")
.finish()
}
}
pub struct BotContextBuilder {
pub url: Option<String>,
pub processors: Vec<Processor>,
}
impl BotContextBuilder {
pub fn new() -> Self {
Self {
url: None,
processors: vec![],
}
}
pub fn url(self, url: impl Into<String>) -> Self {
Self {
url: Some(url.into()),
..self
}
}
pub fn add_processor(
mut self,
processor: impl Into<Processor> + Sync + Send + 'static,
) -> Self {
self.processors.push(processor.into());
self
}
pub fn build(self) -> Result<Arc<BotContext>> {
Ok(Arc::new(BotContext {
connection: Mutex::new(None),
url: self.url,
id: 0,
processors: Arc::new(self.processors),
echo_notifer: Arc::new(DashMap::new()),
}))
}
}
pub struct BotServer {
pub bind: String,
pub processors: Arc<Vec<Processor>>,
}
pub struct BotServerBuilder {
pub bind: Option<String>,
pub processors: Vec<Processor>,
}
impl BotServerBuilder {
pub fn new() -> Self {
Self {
bind: None,
processors: vec![],
}
}
pub fn bind(mut self, bind: impl Into<String>) -> Self {
self.bind = Some(bind.into());
self
}
pub fn add_processor(
mut self,
processor: impl Into<Processor> + Sync + Send + 'static,
) -> Self {
self.processors.push(processor.into());
self
}
pub fn build(self) -> Result<Arc<BotServer>> {
Ok(Arc::new(BotServer {
bind: if let Some(bind) = self.bind {
bind
} else {
return Err(Error::ParamsError("bind must be set".to_string()));
},
processors: Arc::new(self.processors),
}))
}
}
async fn loop_bot<S>(bot_ctx: Arc<BotContext>, ws_stream: WebSocketStream<S>)
where
S: AsyncRead + AsyncWrite + Sync + Send + Unpin + 'static,
{
let (ws_sink, mut split_stream) = ws_stream.split();
let connection = BotConnection {
sender: Box::new(ws_sink),
};
bot_ctx.set_connection(connection).await;
while let Some(msg) = split_stream.next().await {
match msg {
Ok(m) => {
let bot_ctx = bot_ctx.clone();
_ = tokio::spawn(async move { bot_ctx.handle_receive(bot_ctx.clone(), &m).await });
}
Err(e) => {
tracing::error!("WS error: {:?}", e);
break; }
}
}
bot_ctx.set_connection(None).await;
}
pub async fn loop_server(bot_server: Arc<BotServer>) -> Result<()> {
let listener = TcpListener::bind(&bot_server.bind).await.unwrap();
println!("WebSocket server started on ws://{}", &bot_server.bind);
while let Ok((stream, _)) = listener.accept().await {
let processors = bot_server.processors.clone();
tokio::spawn(async move {
let ws_stream = accept_async(stream).await.unwrap();
loop_bot(
Arc::new(BotContext {
connection: Mutex::new(None),
url: None,
id: 0,
processors,
echo_notifer: Arc::new(DashMap::new()),
}),
ws_stream,
)
.await;
});
}
Ok(())
}
pub async fn loop_client(bot_ctx: Arc<BotContext>) -> Result<()> {
let url = bot_ctx
.url
.as_ref()
.ok_or(Error::ParamsError(
"url must be set for loop client".to_string(),
))
.map(|e| e.clone())?;
loop {
match connect_async(&url).await {
Ok((ws_stream, _)) => {
tracing::info!("WS {} Connected!", &url);
let _ = loop_bot(bot_ctx.clone(), ws_stream).await;
}
Err(e) => tracing::error!("WS {} connect error: {:?}", &url, e),
}
tracing::info!("WS {} reconnecting after 15s...", &url);
sleep(Duration::from_secs(15)).await;
}
}