as-rust/exchanges/src/htx_swap_ws.rs

use std::io::Read;
use std::str::from_utf8;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;

use chrono::Utc;
use flate2::bufread::GzDecoder;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use once_cell::sync::Lazy;
use ring::hmac;
use serde_json::{json, Value};
use tokio::sync::Mutex;
use tokio::task;
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::{error, info, trace};

use crate::response_base::ResponseData;
use crate::socket_tool::AbstractWsMode;

pub(crate) static LOGIN_DATA: Lazy<Mutex<(bool, bool)>> = Lazy::new(|| {
    println!("初始化...");
    // 0: 需要登录， 1：是否已经登录
    Mutex::new((false, false))
});


pub enum HtxSwapWsType {
    Public,
    Private,
}


//订阅频道
#[derive(Clone)]
pub enum HtxSwapSubscribeType {
    // 深度
    PuFuturesDepth,
    // // 公开成交
    // PuFuturesDeals,

    // 订单
    PrFuturesOrders,
    // 仓位
    PrFuturesPositions,
    // 余额
    PrFuturesBalances,
}

//账号信息
#[derive(Clone)]
#[allow(dead_code)]
pub struct HtxSwapLogin {
    pub api_key: String,
    pub secret: String,
}

#[derive(Clone)]
pub struct HtxSwapWs {
    //类型
    label: String,
    //地址
    address_url: String,
    //账号信息
    login_param: Option<HtxSwapLogin>,
    //币对
    symbol_s: Vec<String>,
    //订阅
    subscribe_types: Vec<HtxSwapSubscribeType>,
    //心跳间隔
    _heartbeat_time: u64,
}


impl HtxSwapWs {
    /*******************************************************************************************************/
    /*****************************************实例化一个对象****************************************************/
    /*******************************************************************************************************/
    pub fn new(login_param: Option<HtxSwapLogin>, ws_type: HtxSwapWsType) -> HtxSwapWs {
        return HtxSwapWs::new_label("default-HtxSwapWs".to_string(), login_param, ws_type);
    }

    pub fn new_label(label: String, login_param: Option<HtxSwapLogin>, ws_type: HtxSwapWsType) -> HtxSwapWs
    {
        /*******公共频道-私有频道数据组装*/
        let address_url = match ws_type {
            HtxSwapWsType::Public => {
                let url = "wss://api.hbdm.vn/linear-swap-ws".to_string();
                info!("走普通通道(不支持colo通道)：{}", url);
                url
            }
            HtxSwapWsType::Private => {
                let url = "wss://api.hbdm.vn/linear-swap-notification".to_string();
                info!("走普通通道(不支持colo通道)：{}", url);
                url
            }
        };

        HtxSwapWs {
            label,
            address_url,
            login_param,
            symbol_s: vec![],
            subscribe_types: vec![],
            _heartbeat_time: 1000 * 10,
        }
    }

    /*******************************************************************************************************/
    /*****************************************订阅函数********************************************************/
    /*******************************************************************************************************/
    //手动添加订阅信息
    pub fn set_subscribe(&mut self, subscribe_types: Vec<HtxSwapSubscribeType>) {
        self.subscribe_types.extend(subscribe_types);
    }
    //手动添加币对
    pub fn set_symbols(&mut self, mut b_array: Vec<String>) {
        for symbol in b_array.iter_mut() {
            // 大写
            *symbol = symbol.to_uppercase();
            // 字符串替换
            *symbol = symbol.replace("_", "-");
        }
        self.symbol_s = b_array;
    }
    //频道是否需要登录
    fn contains_pr(&self) -> bool {
        for t in self.subscribe_types.clone() {
            if match t {
                HtxSwapSubscribeType::PuFuturesDepth => false,
                // HtxSwapSubscribeType::PuFuturesDeals => false,
                //
                HtxSwapSubscribeType::PrFuturesOrders => true,
                HtxSwapSubscribeType::PrFuturesPositions => true,
                HtxSwapSubscribeType::PrFuturesBalances => true,
            } {
                return true;
            }
        }
        false
    }

    /*******************************************************************************************************/
    /*****************************************工具函数********************************************************/
    /*******************************************************************************************************/
    //订阅枚举解析
    pub fn enum_to_string(symbol: String, subscribe_type: HtxSwapSubscribeType, _login_param: Option<HtxSwapLogin>) -> Value {
        // let access_key;
        // let secret_key;
        // match login_param {
        //     None => {
        //         access_key = "".to_string();
        //         secret_key = "".to_string();
        //     }
        //     Some(param) => {
        //         access_key = param.api_key.clone();
        //         secret_key = param.secret.clone();
        //     }
        // }
        // let cid = "";

        match subscribe_type {
            HtxSwapSubscribeType::PuFuturesDepth => {
                json!({
                      "sub":format!("market.{}.depth.step0", symbol.to_uppercase()),
                      "id": "id5"
                })
            }

            HtxSwapSubscribeType::PrFuturesOrders => {
                json!({
                        "op":"sub",
                        "topic":format!("orders_cross.{}", symbol.to_lowercase())
                })
            }
            HtxSwapSubscribeType::PrFuturesPositions => {
                json!({
                        "op":"sub",
                        "topic":format!("positions_cross.{}", symbol.to_uppercase())
                })
            }
            HtxSwapSubscribeType::PrFuturesBalances => {
                json!({
                        "op":"sub",
                        "topic":format!("accounts_cross.USDT")
                })
            }
        }
    }
    //订阅信息生成
    pub fn get_subscription(&self) -> Vec<Value> {
        let mut args = vec![];
        // 只获取第一个
        // let symbol = self.symbol_s.get(0).unwrap().replace("_", "-");
        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(),
                                                  self.login_param.clone(),
                );
                args.push(ty_str);
            }
        }
        args
    }

    /*******************************************************************************************************/
    /*****************************************socket基本*****************************************************/
    /*******************************************************************************************************/
    //链接
    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(ResponseData) -> Future + Clone + Send + 'static + Sync,
            Future: std::future::Future<Output=()> + Send + 'static, // 确保 Fut 是一个 Future，且输出类型为 ()
    {
        let login_is = self.contains_pr();
        let login_param = self.login_param.clone();
        let subscription = self.get_subscription();
        let address_url = self.address_url.clone();
        let label = self.label.clone();
        // let heartbeat_time = self.heartbeat_time.clone();


        //心跳-- 方法内部线程启动
        // let write_tx_clone1 = Arc::clone(write_tx_am);
        let write_tx_clone2 = Arc::clone(write_tx_am);
        // tokio::spawn(async move {
        //     trace!("线程-异步心跳-开始");
        //     let ping_str = json!({
        //         "method": "server.ping",
        //         "params": {},
        //         "id": 1
        //     });
        //     AbstractWsMode::ping_or_pong(write_tx_clone1, HeartbeatType::Custom(ping_str.to_string()), heartbeat_time).await;
        //     trace!("线程-异步心跳-结束");
        // });

        //设置订阅
        let mut subscribe_array = vec![];


        for s in subscription {
            subscribe_array.push(s.to_string());
        }

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

            info!("启动连接");
            loop {
                info!("htx_usdt_swap socket 连接中……");
                // 需要登录
                if login_is {
                    let mut login_data = LOGIN_DATA.lock().await;
                    let login_param_real = login_param.clone().unwrap();
                    login_data.0 = true;
                    let utc_now = Utc::now();
                    let timestamp = utc_now.format("%Y-%m-%dT%H:%M:%S").to_string();
                    let timestamp_str = percent_encoding::utf8_percent_encode(timestamp.clone().as_str(), crate::htx_swap_rest::FRAGMENT).to_string();

                    let access_key = login_param_real.api_key.clone();
                    let secret_key = login_param_real.secret.clone();
                    let param_str = format!("AccessKeyId={}&SignatureMethod=HmacSHA256&SignatureVersion=2&Timestamp={}", access_key, timestamp_str);

                    let signature = {
                        let message = format!("GET\napi.hbdm.vn\n/linear-swap-notification\n{}", param_str);
                        trace!("组装数据：\n{}", message);

                        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, secret_key.as_bytes());
                        let result = hmac::sign(&hmac_key, &message.as_bytes());
                        let sign = base64::encode(result);
                        sign
                    };
                    let login_param = json!({
                        "op": "auth",
                        "type": "api",
                        "AccessKeyId": access_key,
                        "SignatureMethod": "HmacSHA256",
                        "SignatureVersion": "2",
                        "Timestamp": timestamp,
                        "Signature": signature,
                    });
                    let login_str = login_param.to_string();
                    info!("发起ws登录: {}", login_str);
                    let write_tx_c = Arc::clone(&write_tx_clone2);
                    AbstractWsMode::send_subscribe(write_tx_c, Message::Text(login_str)).await;
                }

                AbstractWsMode::ws_connect_async(is_shutdown_arc.clone(), handle_function.clone(), address_url.clone(),
                                                 login_is, label.clone(), subscribe_array.clone(), write_to_socket_rx_arc.clone(),
                                                 Self::message_text_sync, Self::message_ping, Self::message_pong, Self::message_binary_sync,
                ).await;
                let mut login_data = LOGIN_DATA.lock().await;
                // 断联后 设置为没有登录
                login_data.1 = false;
                info!("htx_usdt_swap socket 断连，1s以后重连……");
                error!("htx_usdt_swap socket 断连，1s以后重连……");
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        });
        tokio::try_join!(t2).unwrap();
        trace!("线程-心跳与链接-结束");

        Ok(())
    }
    /*******************************************************************************************************/
    /*****************************************数据解析*****************************************************/
    /*******************************************************************************************************/
    //数据解析-Text
    pub async fn message_text(text: String) -> Option<ResponseData> {
        let response_data = Self::ok_text(text).await;
        Option::from(response_data)
    }
    pub fn message_text_sync(text: String) -> Option<ResponseData> {
        // 使用 tokio::task::block_in_place 来等待异步函数的结果
        task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(Self::message_text(text))
        })
    }
    //数据解析-ping
    pub fn message_ping(_pi: Vec<u8>) -> Option<ResponseData> {
        return Option::from(ResponseData::new("".to_string(), -300, "success".to_string(), Value::Null));
    }
    //数据解析-pong
    pub fn message_pong(_po: Vec<u8>) -> Option<ResponseData> {
        return Option::from(ResponseData::new("".to_string(), -301, "success".to_string(), Value::Null));
    }
    //数据解析-二进制
    pub async fn message_binary(binary: Vec<u8>) -> Option<ResponseData> {
        //二进制WebSocket消息
        let message_str = Self::parse_zip_data(binary);
        let response_data = Self::ok_text(message_str).await;
        Option::from(response_data)
    }
    pub fn message_binary_sync(binary: Vec<u8>) -> Option<ResponseData> {
        // 使用 tokio::task::block_in_place 来等待异步函数的结果
        task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(Self::message_binary(binary))
        })
    }
    //数据解析
    pub async fn ok_text(text: String) -> ResponseData
    {
        trace!("原始数据：{}", text);
        let mut res_data = ResponseData::new("".to_string(), 200, "success".to_string(), Value::Null);
        let json_value: Value = serde_json::from_str(&text).unwrap();

        /*公共：响应*/
        //心跳包
        let ping = json_value["ping"].as_i64();
        match ping {
            Some(ping_ts) => {
                let pong = json!({
                    "pong": ping_ts
                });
                return ResponseData::new("".to_string(), -302, "success".to_string(), pong);
            }
            None => {}
        }
        //推送数据
        let status = json_value["status"].as_str();
        match status {
            Some(v) => {
                match v {
                    "ok" => {
                        res_data.channel = format!("{}", v);
                        res_data.code = -201;
                        res_data.data = json_value["data"].clone();
                    }
                    "error" => {
                        res_data.code = 400;
                        res_data.message = format!("{}", json_value["err-msg"].as_str().unwrap());
                        res_data.channel = format!("{}", json_value["id"].as_str().unwrap());
                    }
                    &_ => {}
                }
                return res_data;
            }
            None => {}
        }
        let ch = json_value["ch"].as_str();
        match ch {
            Some(channel) => {
                res_data.channel = channel.parse().unwrap();
                res_data.code = 200;
                res_data.data = json_value["tick"].clone();
                return res_data;
            }
            None => {}
        }


        /*私有：响应*/
        let op = json_value["op"].as_str().unwrap();
        match op {
            "auth" => {
                let op = json_value["err-code"].as_i64().unwrap();
                res_data.channel = "auth".to_string();
                if op == 0 {
                    res_data.code = -200;
                    res_data.message = "登录成功".to_string();
                    res_data.data = json_value["data"].clone();
                } else {
                    res_data.code = 400;
                    res_data.message = format!("登录失败:{}", json_value["err-msg"].as_str().unwrap());
                }
                return res_data;
            }
            "ping" => {
                let ts = json_value["ts"].as_str().unwrap();
                let pong = json!({
                    "op": "pong",
                    "ts": ts
                });
                return ResponseData::new("".to_string(), -302, "success".to_string(), pong);
            }
            "sub" => {
                res_data.channel = json_value["topic"].as_str().unwrap().to_string();
                res_data.code = -201;
                res_data.message = "订阅成功".to_string();
                return res_data;
            }
            "notify" => {
                res_data.channel = json_value["topic"].as_str().unwrap().to_string();
                res_data.code = 200;
                res_data.message = "推送数据".to_string();
                if json_value.get("data").is_some() {
                    res_data.data = json_value["data"].clone();
                } else {
                    res_data.data = json_value.clone();
                }

                return res_data;
            }
            _ => {}
        }


        res_data.code = 400;
        res_data.message = format!("未知响应内容");
        res_data.data = text.parse().unwrap();
        trace!("--------------------------------");
        res_data
    }

    fn parse_zip_data(p0: Vec<u8>) -> String {
        // 创建一个GzDecoder的实例，将压缩数据作为输入
        let mut decoder = GzDecoder::new(&p0[..]);

        // 创建一个缓冲区来存放解压缩后的数据
        let mut decompressed_data = Vec::new();

        // 读取解压缩的数据到缓冲区中
        decoder.read_to_end(&mut decompressed_data).expect("解压缩失败");
        let result = from_utf8(&decompressed_data)
            .expect("解压缩后的数据不是有效的UTF-8");

        // info!("解压缩数据 {:?}", result);
        result.to_string()
    }
}