as-rust/exchanges/src/okx_swap_ws.rs

use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;

use chrono::Utc;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use ring::hmac;
use serde_json::{json, Value};
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::{info, trace};

use crate::response_base::ResponseData;
use crate::socket_tool::{AbstractWsMode, HeartbeatType};
use crate::utils::get_time_microsecond;

//类型
pub enum OkxSwapWsType {
    //订阅频道类型
    Public,
    Private,
    Business,
}

//订阅频道
#[derive(Clone)]
pub enum OkxSwapSubscribeType {
    PuIndexTickers,
    PuBooks5,
    Putrades,
    PuBooks50L2tbt,

    BuIndexCandle30m,

    PrBalanceAndPosition,
    PrAccount(String),
    PrOrders,
    PrPositions,

}

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

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

impl OkxSwapWs {
    /*******************************************************************************************************/
    /*****************************************获取一个对象****************************************************/
    /*******************************************************************************************************/
    pub fn new(is_colo: bool, login_param: Option<OkxSwapLogin>, ws_type: OkxSwapWsType) -> OkxSwapWs {
        return OkxSwapWs::new_label("default-OkxSwapWs".to_string(), is_colo, login_param, ws_type);
    }
    pub fn new_label(label: String, is_colo: bool, login_param: Option<OkxSwapLogin>, ws_type: OkxSwapWsType) -> OkxSwapWs {
        /*******公共频道-私有频道数据组装*/
        let address_url = match ws_type {
            OkxSwapWsType::Public => {
                "wss://ws.okx.com:8443/ws/v5/public".to_string()
            }
            OkxSwapWsType::Private => {
                "wss://ws.okx.com:8443/ws/v5/private".to_string()
            }
            OkxSwapWsType::Business => {
                "wss://ws.okx.com:8443/ws/v5/business".to_string()
            }
        };

        if is_colo {
            info!("开启高速(未配置，走普通：{})通道",address_url);
        } else {
            info!("走普通通道：{}",address_url);
        }
        /*****返回结构体*******/
        OkxSwapWs {
            label,
            address_url,
            login_param,
            symbol_s: vec![],
            subscribe_types: vec![],
            heartbeat_time: 1000 * 5,
        }
    }

    /*******************************************************************************************************/
    /*****************************************订阅函数********************************************************/
    /*******************************************************************************************************/
    //手动添加订阅信息
    pub fn set_subscribe(&mut self, subscribe_types: Vec<OkxSwapSubscribeType>) {
        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 {
                OkxSwapSubscribeType::PuIndexTickers => false,
                OkxSwapSubscribeType::PuBooks5 => false,
                OkxSwapSubscribeType::Putrades => false,
                OkxSwapSubscribeType::PuBooks50L2tbt => false,

                OkxSwapSubscribeType::BuIndexCandle30m => false,

                OkxSwapSubscribeType::PrBalanceAndPosition => true,
                OkxSwapSubscribeType::PrAccount(_) => true,
                OkxSwapSubscribeType::PrOrders => true,
                OkxSwapSubscribeType::PrPositions => true,
            } {
                return true;
            }
        }
        false
    }
    /*******************************************************************************************************/
    /*****************************************工具函数********************************************************/
    /*******************************************************************************************************/
    //订阅枚举解析
    pub fn enum_to_string(symbol: String, subscribe_type: OkxSwapSubscribeType) -> Value {
        match subscribe_type {
            OkxSwapSubscribeType::PuIndexTickers => {
                json!({
                    "channel":"index-tickers",
                    "instId":symbol
                })
            }

            OkxSwapSubscribeType::PuBooks5 => {
                json!({
                    "channel":"books5",
                    "instId":symbol
                })
            }
            OkxSwapSubscribeType::Putrades => {
                json!({
                    "channel":"trades",
                    "instId":symbol
                })
            }

            OkxSwapSubscribeType::BuIndexCandle30m => {
                json!({
                    "channel":"index-candle30m",
                    "instId":symbol
                })
            }

            OkxSwapSubscribeType::PrAccount(ccy) => {
                json!({
                    "channel":"account",
                    "ccy":ccy
                })
            }
            OkxSwapSubscribeType::PuBooks50L2tbt => {
                json!({
                    "channel":"books50-l2-tbt",
                    "instId":symbol
                })
            }
            OkxSwapSubscribeType::PrBalanceAndPosition => {
                json!({
                    "channel":"balance_and_position"
                })
            }
            OkxSwapSubscribeType::PrOrders => {
                json!({
                    "channel":"orders",
                    "instType":"SWAP",
                    "instFamily":symbol
                })
            }
            OkxSwapSubscribeType::PrPositions => {
                json!({
                    "channel":"positions",
                    "instType":"SWAP",
                })
            }
        }
    }
    //订阅信息生成
    pub fn get_subscription(&self) -> String {
        let mut args = 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());
                args.push(ty_str);
            }
        }
        let str = json!({
            "op": "subscribe",
            "args": args
        });

        // trace!("订阅信息：{}", str.to_string());

        str.to_string()
    }
    //登录组装
    fn log_in_to_str(login_param: Option<OkxSwapLogin>) -> String {
        let mut login_json_str = "".to_string();

        let mut access_key: String = "".to_string();
        let mut secret_key: String = "".to_string();
        let mut passphrase: String = "".to_string();


        if let Some(param) = login_param {
            access_key = param.api_key;
            secret_key = param.secret_key;
            passphrase = param.passphrase;
        }

        if access_key.len() > 0 || secret_key.len() > 0 || passphrase.len() > 0 {
            let timestamp = Utc::now().timestamp().to_string();
            // 时间戳 + 请求类型+ 请求参数字符串
            let message = format!("{}GET{}", timestamp, "/users/self/verify");
            let hmac_key = ring::hmac::Key::new(hmac::HMAC_SHA256, secret_key.as_bytes());
            let result = ring::hmac::sign(&hmac_key, &message.as_bytes());
            let sign = base64::encode(result);

            let login_json = json!({
                              "op": "login",
                                "args": [{
                                "apiKey": access_key,
                                "passphrase": passphrase,
                                "timestamp": timestamp,
                                "sign": sign  }]
                        });

            trace!("---login_json：{0}", login_json.to_string());
            trace!("--登陆：{}", login_json.to_string());
            login_json_str = login_json.to_string();
        }
        login_json_str
    }
    /*******************************************************************************************************/
    /*****************************************socket基本*****************************************************/
    /*******************************************************************************************************/
    //链接
    pub async fn ws_connect_async(&mut self,
                                  bool_v1: Arc<AtomicBool>,
                                  write_tx_am: &Arc<Mutex<UnboundedSender<Message>>>,
                                  write_rx: UnboundedReceiver<Message>,
                                  read_tx: UnboundedSender<ResponseData>) -> Result<(), Error>
    {
        let login_is = self.contains_pr();
        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);
        tokio::spawn(async move {
            trace!("线程-异步心跳-开始");
            AbstractWsMode::ping_or_pong(write_tx_clone1, HeartbeatType::Ping, heartbeat_time).await;
            trace!("线程-异步心跳-结束");
        });

        //设置订阅
        let  subscribe_array = vec![];
        if login_is {
            let write_tx_clone2 = Arc::clone(write_tx_am);
            let login_str = Self::log_in_to_str(self.login_param.clone());
            tokio::spawn(async move {
                //登录相关
                AbstractWsMode::send_subscribe(write_tx_clone2, Message::Text(login_str)).await;
            });
        }
        let write_tx_clone3 = Arc::clone(write_tx_am);
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(3 * 1000)).await;
            //登录相关
            AbstractWsMode::send_subscribe(write_tx_clone3, Message::Text(subscription)).await;
        });

        //链接
        let t2 = tokio::spawn(async move {
            trace!("线程-异步链接-开始");
            AbstractWsMode::ws_connect_async(bool_v1, address_url.clone(),
                                             label.clone(), subscribe_array,
                                             write_rx, read_tx,
                                             Self::message_text,
                                             Self::message_ping,
                                             Self::message_pong,
            ).await.expect("OKX-期货");
            trace!("线程-异步链接-结束");
        });
        tokio::try_join!(t2).unwrap();
        trace!("线程-心跳与链接-结束");

        Ok(())
    }

    // async fn run(&mut self, bool_v1: Arc<AtomicBool>)
    // {
    //     //订阅信息组装
    //     let subscription = self.get_subscription();
    //     loop {
    //         trace!("要连接咯~~！！{}", self.address_url);
    //
    //         let address_url = Url::parse(self.address_url.as_str()).unwrap();
    //         //1. 判断是否需要代理,根据代理地址是否存来选择
    //         if self.proxy.ip_address.len() > 0 {
    //             let ip_array: Vec<&str> = self.proxy.ip_address.split(".").collect();
    //             let proxy_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(
    //                 ip_array[0].parse().unwrap(),
    //                 ip_array[1].parse().unwrap(),
    //                 ip_array[2].parse().unwrap(),
    //                 ip_array[3].parse().unwrap())
    //             ), self.proxy.port.parse().unwrap());
    //             let websocket_config = Some(WebSocketConfig {
    //                 max_send_queue: Some(16),
    //                 max_message_size: Some(16 * 1024 * 1024),
    //                 max_frame_size: Some(16 * 1024 * 1024),
    //                 accept_unmasked_frames: false,
    //             });
    //             let max_redirects = 5;
    //             match connect_with_proxy(address_url.clone(),
    //                                      proxy_address, websocket_config, max_redirects) {
    //                 Ok(ws) => {
    //                     let bool_v1_clone = Arc::clone(&bool_v1);
    //                     self.proxy_subscription(bool_v1_clone, ws.0, subscription.clone()).await;
    //                 }
    //                 Err(err) => {
    //                     trace!("Can't connect(无法连接): {}", err);
    //                 }
    //             };
    //         } else {
    //             match connect(address_url.clone()) {
    //                 Ok(ws) => {
    //                     let bool_v1_clone = Arc::clone(&bool_v1);
    //                     self.subscription(bool_v1_clone, ws.0, subscription.clone()).await;
    //                 }
    //                 Err(err) => {
    //                     // 连接失败时执行的操作
    //                     trace!("Can't connect(无法连接): {}", err);
    //                     // 返回一个默认的 WebSocket 对象或其他适当的值
    //                     // 或者根据需要触发 panic 或返回错误信息
    //                 }
    //             };
    //         }
    //         trace!("退出来咯");
    //
    //         let bool_v1_clone = Arc::clone(&bool_v1);
    //         let bool_v1_v = bool_v1_clone.load(Ordering::SeqCst);
    //         if !bool_v1_v {
    //             break;
    //         }
    //     }
    // }
    //
    // //代理
    // async fn proxy_subscription(&self, bool_v1: Arc<AtomicBool>, mut web_socket: WebSocket<ProxyAutoStream>,
    //                             subscription: String)
    // {
    //     let lable = self.label.clone();
    //     /*****登陆***/
    //     let login_str = self.log_in_to_str();
    //     if login_str != "" {
    //         let _ = web_socket.write_message(Message::Text(login_str));
    //         thread::sleep(Duration::from_secs(3));
    //     }
    //     /*****订阅***/
    //     web_socket.write_message(Message::Text(subscription))
    //         .unwrap();
    //     /*****消息溜***/
    //     let mut ping_timeout = chrono::Utc::now().timestamp_millis();
    //     loop {
    //         tokio::time::sleep(Duration::from_millis(1)).await;
    //         let msg = web_socket.read_message();
    //         match msg {
    //             Ok(Message::Text(text)) => {
    //                 let get_time = chrono::Utc::now().timestamp_millis();
    //                 if (get_time - ping_timeout) >= (1000 * 30) {
    //                     trace!("30s 一次主动发送心跳包!");
    //                     let _ = web_socket.write_message(Message::Ping(Vec::from("ping")));
    //                     ping_timeout = get_time;
    //                 }
    //
    //                 let mut res_data = Self::ok_text(lable.to_string(), text);
    //                 res_data.time = get_time_microsecond();
    //                 if res_data.code == "-201" {
    //                     trace!("登陆成功!");
    //                 } else if res_data.code == "-200" {
    //                     trace!("订阅成功：{:?}", res_data.data);
    //                 } else {
    //                     self.sender.send(res_data).await.unwrap();
    //                 }
    //             }
    //             Ok(Message::Ping(s)) => {
    //                 trace!( "Ping-响应--{:?}", String::from_utf8(s.clone()));
    //                 let _ = web_socket.write_message(Message::Pong(Vec::from("pong")));
    //                 trace!( "回应-pong---{:?}", String::from_utf8(s.clone()));
    //             }
    //             Ok(Message::Pong(s)) => {
    //                 // trace!("Pong-响应--{:?}", String::from_utf8(s));
    //                 trace!( "Pong-响应--{:?}", String::from_utf8(s.clone()));
    //             }
    //             Ok(Message::Close(_)) => {
    //                 // trace!("socket 关闭: ");
    //                 trace!( "Close-响应");
    //             }
    //             Err(e) => {
    //                 // trace!("Error receiving message: {}", error);
    //                 trace!("Err-响应{}", e);
    //                 error!( "Err-响应{}", e);
    //                 break;
    //             }
    //             _ => {}
    //         }
    //
    //         let bool_v1_v = bool_v1.load(Ordering::SeqCst);
    //         if !bool_v1_v {
    //             break;
    //         }
    //     }
    //     web_socket.close(None).unwrap();
    // }
    //
    // //非代理
    // async fn subscription(&self, bool_v1: Arc<AtomicBool>, mut web_socket: WebSocket<AutoStream>,
    //                       subscription: String)
    // {
    //     let lable = self.label.clone();
    //     /*****订阅***/
    //     web_socket.write_message(Message::Text(subscription))
    //         .unwrap();
    //     /*****消息溜***/
    //     let mut ping_timeout = chrono::Utc::now().timestamp_millis();
    //     loop {
    //         tokio::time::sleep(Duration::from_millis(1)).await;
    //         let msg = web_socket.read_message();
    //         match msg {
    //             Ok(Message::Text(text)) => {
    //                 let get_time = chrono::Utc::now().timestamp_millis();
    //                 if (get_time - ping_timeout) >= (1000 * 30) {
    //                     trace!("30s 一次主动发送心跳包!");
    //                     let _ = web_socket.write_message(Message::Ping(Vec::from("ping")));
    //                     ping_timeout = get_time;
    //                 }
    //
    //                 let mut res_data = Self::ok_text(lable.to_string(), text);
    //                 res_data.time = get_time_microsecond();
    //                 if res_data.code == "-201" {
    //                     trace!("登陆成功!");
    //                 } else if res_data.code == "-200" {
    //                     trace!("订阅成功：{:?}", res_data.data);
    //                 } else {
    //                     self.sender.send(res_data).await.unwrap();
    //                 }
    //             }
    //             Ok(Message::Ping(s)) => {
    //                 trace!( "Ping-响应--{:?}", String::from_utf8(s.clone()));
    //                 let _ = web_socket.write_message(Message::Pong(Vec::from("pong")));
    //                 trace!( "回应-pong---{:?}", String::from_utf8(s.clone()));
    //             }
    //             Ok(Message::Pong(s)) => {
    //                 // trace!("Pong-响应--{:?}", String::from_utf8(s));
    //                 trace!( "Pong-响应--{:?}", String::from_utf8(s.clone()));
    //             }
    //             Ok(Message::Close(_)) => {
    //                 // trace!("socket 关闭: ");
    //                 trace!( "Close-响应");
    //             }
    //             Err(e) => {
    //                 // trace!("Error receiving message: {}", error);
    //                 trace!("Err-响应{}", e);
    //                 error!( "Err-响应{}", e);
    //                 break;
    //             }
    //             _ => {}
    //         }
    //
    //         let bool_v1_v = bool_v1.load(Ordering::SeqCst);
    //         if !bool_v1_v {
    //             break;
    //         }
    //     }
    //     web_socket.close(None).unwrap();
    // }
    /*******************************************************************************************************/
    /*****************************************数据解析*****************************************************/
    /*******************************************************************************************************/
    //数据解析-Text
    pub fn message_text(text: String) -> Option<ResponseData> {
        let mut response_data = Self::ok_text(text);
        response_data.time = get_time_microsecond();
        match response_data.code.as_str() {
            "200" => Option::from(response_data),
            "-200" => {
                trace!("订阅成功：{:?}", response_data);
                None
            }
            "-201" => {
                trace!("登录成功：{:?}", response_data);
                None
            }
            _ => None
        }
    }
    //数据解析-ping
    pub fn message_ping(_pi: Vec<u8>) -> Option<ResponseData> {
        trace!("服务器响应-ping");
        return None;
    }
    //数据解析-pong
    pub fn message_pong(_po: Vec<u8>) -> Option<ResponseData> {
        trace!("服务器响应-pong");
        return None;
    }
    //数据解析
    pub fn ok_text(text: String) -> ResponseData
    {
        // trace!("元数据:{}",text);
        let mut res_data = ResponseData::new("".to_string(), "200".to_string(), "success".to_string(), "".to_string());
        let json_value: serde_json::Value = serde_json::from_str(&text).unwrap();
        if json_value.get("event").is_some() {//订阅返回
            if json_value["event"].as_str() == Option::from("login") &&
                json_value["code"].as_str() == Option::from("0") {
                res_data.code = "-201".to_string();
                res_data.message = format!("登陆成功!");
            } else if json_value["event"].as_str() == Option::from("error") {
                res_data.code = json_value["code"].to_string();
                res_data.message = format!("订阅失败：{}", json_value["msg"].to_string());
            } else if json_value["event"].as_str() == Option::from("subscribe") {
                res_data.code = "-200".to_string();
                res_data.data = text;
                res_data.message = format!("订阅成功!");
            } else {
                res_data.data = text;
            }
        } else {
            if json_value.get("arg").is_some() && json_value.get("data").is_some() {
                res_data.channel = format!("{}", json_value["arg"]["channel"].as_str().unwrap());
                res_data.data = json_value["data"].to_string();
                // res_data.reach_time = json_value["data"][0]["ts"].as_str().unwrap().parse().unwrap()
            } else {
                res_data.data = text;
                res_data.channel = "未知频道".to_string();
            }
        }
        res_data
    }
}