pin_trading_tool/src/exchange/mexc_spot_ws.rs

use std::fmt::format;
use std::io::Read;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;

use flate2::read::GzDecoder;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use serde_json::json;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::{error, info, trace};

use crate::exchange::response_base::Response;
use crate::exchange::socket_tool::AbstractWsMode;

//类型
pub enum MexcSpotWsType {
    PublicAndPrivate,
}


#[derive(Debug)]
#[derive(Clone)]
pub struct MexcSpotWsParam {
    pub token: String,
    pub ws_url: String,
    pub ws_ping_interval: i64,
    pub ws_ping_timeout: i64,
    pub is_ok_subscribe: bool,
}

//订阅频道
#[derive(Clone)]
pub enum MexcSpotWsSubscribeType {
    // 深度
    PuFuturesDepth,
    // K线数据，Min -> 分钟; Hour -> 小时; Day -> 天; Week -> 周, M -> 月
    // Min1
    // Min5
    // Min15
    // Min30
    // Min60
    // Hour4
    // Hour8
    // Day1
    // Week1
    // Month1
    PuFuturesRecords(String),
}

//账号信息
#[derive(Clone, Debug)]
pub struct MexcSpotWsLogin {
    pub access_key: String,
    pub secret_key: String,
    pub pass_key: String,
}

#[derive(Clone)]
#[allow(dead_code)]
pub struct MexcSpotWs {
    //类型
    tag: String,
    //地址
    address_url: String,
    //账号
    login_param: Option<MexcSpotWsLogin>,
    //登录数据
    ws_param: MexcSpotWsParam,
    //币对
    symbol_s: Vec<String>,
    //订阅
    subscribe_types: Vec<MexcSpotWsSubscribeType>,
    //心跳间隔
    heartbeat_time: u64,
}

impl MexcSpotWs {
    // ============================================= 构造函数 ================================================
    pub fn new_with_tag(tag: String, login_param: Option<MexcSpotWsLogin>, ws_type: MexcSpotWsType) -> MexcSpotWs {
        /*******公共频道-私有频道数据组装*/
        let address_url = match ws_type {
            MexcSpotWsType::PublicAndPrivate => {
                let url = "ws://wbs-api.mexc.com/ws".to_string();
                url
            }
        };

        /*******公共频道-私有频道数据组装*/
        let ws_param = MexcSpotWsParam {
            token: "".to_string(),
            ws_url: "".to_string(),
            ws_ping_interval: 0,
            ws_ping_timeout: 0,
            is_ok_subscribe: false,
        };

        MexcSpotWs {
            tag,
            address_url,
            login_param,
            ws_param,
            symbol_s: vec![],
            subscribe_types: vec![],
            heartbeat_time: 1000 * 18,
        }
    }

    // ============================================= 订阅函数 ================================================
    // 手动添加订阅信息
    pub fn set_subscribe(&mut self, subscribe_types: Vec<MexcSpotWsSubscribeType>) {
        self.subscribe_types.extend(subscribe_types);
    }
    // 手动添加币对
    pub fn set_symbols(&mut self, mut symbol_array: Vec<String>) {
        for symbol in symbol_array.iter_mut() {
            // 大写
            *symbol = symbol.to_uppercase();
            // 字符串替换
            *symbol = symbol.replace("_", "_");
        }
        self.symbol_s = symbol_array;
    }
    fn contains_pr(&self) -> bool {
        for t in self.subscribe_types.clone() {
            if match t {
                MexcSpotWsSubscribeType::PuFuturesRecords(_) => false,
                MexcSpotWsSubscribeType::PuFuturesDepth => false,
            } {
                return true;
            }
        }
        false
    }

    // 订阅枚举解析
    pub fn enum_to_string(symbol: String, subscribe_type: MexcSpotWsSubscribeType) -> Value {
        match subscribe_type {
            // 深度
            MexcSpotWsSubscribeType::PuFuturesDepth => {
                json!({
                    "method": "SUBSCRIPTION",
                    "param": [
                        format!("spot@public.aggre.depth.v3.api.pb@10ms@{symbol}")
                    ]
                })
            }
            // k线
            MexcSpotWsSubscribeType::PuFuturesRecords(interval) => {
                json!({
                    "method": "SUBSCRIPTION",
                    "param": [
                        format!("spot@public.kline.v3.api.pb@{symbol}@{interval}")
                    ]
                })
            }
        }
    }
    // 订阅信息生成
    pub fn get_subscription(&self) -> Vec<String> {
        let mut array = vec![];
        for symbol in &self.symbol_s {
            for subscribe_type in &self.subscribe_types {
                let ty_str = Self::enum_to_string(symbol.clone(), subscribe_type.clone());
                array.push(ty_str.to_string());
            }
        }
        array
    }

    // 链接
    pub async fn ws_connect_async<F, Future>(&mut self,
                                             is_shutdown_arc: Arc<AtomicBool>,
                                             handle_function: F,
                                             _write_tx_am: &Arc<Mutex<UnboundedSender<Message>>>,
                                             write_to_socket_rx: UnboundedReceiver<Message>) -> Result<(), Error>
        where
            F: Fn(Response) -> Future + Clone + Send + 'static + Sync,
            Future: std::future::Future<Output=()> + Send + 'static, // 确保 Fut 是一个 Future，且输出类型为 ()
    {
        let login_is = self.contains_pr();
        let subscription = self.get_subscription();
        let address_url = self.address_url.clone();
        let tag = self.tag.clone();
        // let heartbeat_time = self.ws_param.ws_ping_interval.clone();

        //心跳-- 方法内部线程启动
        // let write_tx_clone1 = write_tx_am.clone();
        // tokio::spawn(async move {
        //     trace!("线程-异步心跳-开始");
        //     AbstractWsMode::ping_or_pong(write_tx_clone1, HeartbeatType::Ping, heartbeat_time as u64).await;
        //     trace!("线程-异步心跳-结束");
        // });


        // 设置订阅
        let subscribe_array = subscription.clone();
        if login_is {
            //登录相关
        }

        // 链接
        let t2 = tokio::spawn(async move {
            let write_to_socket_rx_arc = Arc::new(Mutex::new(write_to_socket_rx));

            loop {
                info!("Mexc_usdt_swap socket 连接中……");
                AbstractWsMode::ws_connect_async(is_shutdown_arc.clone(), handle_function.clone(), address_url.clone(),
                                                 false, tag.clone(), subscribe_array.clone(), write_to_socket_rx_arc.clone(),
                                                 Self::message_text, Self::message_ping, Self::message_pong, Self::message_binary).await;

                error!("Mexc_usdt_swap socket 断连，1s以后重连……");
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        });

        tokio::try_join!(t2).unwrap();
        trace!("线程-心跳与链接-结束");

        Ok(())
    }

    //数据解析-Text
    pub fn message_text(text: String) -> Option<Response> {
        let response_data = Self::ok_text(text);
        Option::from(response_data)
    }
    //数据解析-ping
    pub fn message_ping(_pi: Vec<u8>) -> Option<Response> {
        Option::from(Response::new("".to_string(), -300, "success".to_string(), Value::Null))
    }
    //数据解析-pong
    pub fn message_pong(_po: Vec<u8>) -> Option<Response> {
        Option::from(Response::new("".to_string(), -301, "success".to_string(), Value::Null))
    }
    //数据解析-二进制
    pub fn message_binary(po: Vec<u8>) -> Option<Response> {
        //二进制WebSocket消息
        // let message_str = format!("Binary:{:?}", _po);
        // Option::from(ResponseData::new("".to_string(), 2, message_str, Value::Null))
        // let result = String::from_utf8(bytes);
        // let result = String::from_utf8(po);

        let mut gz_decoder = GzDecoder::new(&po[..]);
        let mut decompressed_data = Vec::new();

        // 尝试解压数据
        if let Ok(_) = gz_decoder.read_to_end(&mut decompressed_data) {
            // 将解压后的字节向量转换为 UTF-8 字符串
            match String::from_utf8(decompressed_data) {
                Ok(text) => {
                    let response_data = Self::ok_text(text);
                    return Option::from(response_data);
                }
                Err(_) => {
                    return Option::from(Response::new("".to_string(), 400, "二进制数据转化出错".to_string(), Value::Null));
                }
            }
        } else {
            return Option::from(Response::new("".to_string(), 400, "二进制数据转化出错".to_string(), Value::Null));
        }
    }
    //数据解析
    pub fn ok_text(text: String) -> Response
    {
        let mut res_data = Response::new("".to_string(), 200, "success".to_string(), Value::Null);
        let json_value: Value = serde_json::from_str(&text).unwrap();

        match json_value["channel"].as_str() {
            Some(method) => {
                if method.contains("pong") {
                    return Response::new("".to_string(), -301, "success".to_string(), Value::Null);
                } else if method.contains("rs.sub.") {
                    //订阅响应
                    let data = json_value["data"].as_str().unwrap();
                    if method.contains(".depth") {
                        res_data.channel = "futures.order_book".to_string();
                    } else if method.contains(".kline") {
                        res_data.channel = "futures.candlesticks".to_string();
                    } else if method.contains(".deal") {
                        res_data.channel = "futures.trades".to_string();
                    } else {
                        res_data.channel = "未知频道订阅".to_string();
                    }

                    if data == "success" {
                        res_data.code = -201;
                        res_data.message = "订阅成功".to_string();
                    } else {
                        res_data.code = 400;
                        res_data.message = "订阅失败".to_string();
                    }
                } else if method.contains("push.") {
                    if method.contains(".depth") {
                        res_data.channel = "futures.order_book".to_string();
                    } else if method.contains(".kline") {
                        res_data.channel = "futures.candlesticks".to_string();
                    } else {
                        res_data.channel = "未知频道推送".to_string();
                    }
                    res_data.code = 200;
                    res_data.data = json_value.clone();
                } else {
                    res_data.code = -1;
                    res_data.message = "未知解析".to_string();
                }
            }
            None => {
                res_data.code = -1;
                res_data.message = "未知解析".to_string();
            }
        }

        res_data
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::AtomicBool;
    use tokio::sync::Mutex;
    use tokio_tungstenite::tungstenite::Message;
    use tracing::info;
    use crate::exchange::mexc_spot_ws::{MexcSpotWs, MexcSpotWsSubscribeType, MexcSpotWsType};
    use crate::exchange::response_base::Response;
    use crate::utils::log_setup::setup_logging;

    #[tokio::test]
    async fn test_mexc_spot_ws() {
        let ws_running = Arc::new(AtomicBool::new(true));
        let (write_tx, write_rx) = futures_channel::mpsc::unbounded::<Message>();
        let _guard = setup_logging().unwrap();

        let mut ws = MexcSpotWs::new_with_tag("Mexc".to_string(), None, MexcSpotWsType::PublicAndPrivate);

        ws.set_subscribe(vec![
            MexcSpotWsSubscribeType::PuFuturesRecords("Min1".to_string())
        ]);

        ws.set_symbols(vec!["BTCUSDT".to_string()]);

        let fun = move |response: Response| {
            info!("{}", serde_json::to_string_pretty(&response.data).unwrap());

            async move {}
        };

        // 链接
        info!("开始链接");
        let write_tx_am = Arc::new(Mutex::new(write_tx));
        ws.ws_connect_async(ws_running, fun, &write_tx_am, write_rx)
            .await
            .expect("链接失败");
    }
}