pin_trading_tool/src/exchange/mexc_spot_ws.rs

use std::fmt::format;
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;

use flate2::read::GzDecoder;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use prost::Message as ProstMessage;
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,
}

pub mod mexc_spot {
    include!(concat!(env!("OUT_DIR"), "/_.rs"));
}

use mexc_spot::PublicSpotKlineV3ApiMessage;
use mexc_spot::KlineDataV3;
// 引入新的结构体
use mexc_spot::PublicIncreaseDepthsV3ApiMessage;
use mexc_spot::DepthDataContentV3;
use mexc_spot::DepthItemV3;

#[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 = "wss://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",
                    "params": [
                        format!("spot@public.aggre.depth.v3.api.pb@10ms@{symbol}")
                    ]
                })
            }
            // k线
            MexcSpotWsSubscribeType::PuFuturesRecords(interval) => {
                json!({
                    "method": "SUBSCRIPTION",
                    "params": [
                        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> {
        // info!("Received binary message ({} bytes)", po.len());

        // 1. 尝试用新的顶层消息结构 PublicSpotKlineV3ApiMessage 来解析 K 线数据
        // 根据 Topic 前缀判断依然有效，但现在是判断是否**可能**是 K 线相关消息
        let prefix_len = po.len().min(100);
        let prefix_string = String::from_utf8_lossy(&po[..prefix_len]);

        if prefix_string.contains("spot@public.kline.v3.api.pb") {
            // info!("通过 Topic 前缀判断为 K 线数据相关消息");

            // 尝试解析为 PublicSpotKlineV3ApiMessage
            match PublicSpotKlineV3ApiMessage::decode(&po[..]) {
                Ok(kline_message) => {
                    // info!("成功解析为顶层 K 线消息结构");
                    // 检查是否包含嵌套的 KlineDataV3 字段 (Tag 308)
                    if let Some(kline_data) = kline_message.kline_data { // 注意这里 PublicSpotKlineV3ApiMessage 的 kline_data 字段是 Option<KlineDataV3>
                        // info!("找到并成功访问嵌套的 KlineDataV3");
                        // 现在 kline_data 是 KlineDataV3 结构体，你可以使用它了！
                        // 填充 Response 并返回 (省略详细实现)
                        let response_data = Response::new(
                            kline_message.topic_info.clone(), // 使用解析到的 Topic 信息
                            200,
                            "OK".to_string(),
                        json!({
                                "interval": kline_data.interval,
                                "windowStart": kline_data.window_start, //注意 snake_case
                                "openingPrice": kline_data.opening_price,
                                "closingPrice": kline_data.closing_price,
                                "highestPrice": kline_data.highest_price,
                                "lowestPrice": kline_data.lowest_price,
                                "volume": kline_data.volume,
                                "amount": kline_data.amount,
                                "windowEnd": kline_data.window_end,
                                // 可以添加顶层字段的信息，如果需要
                                "topic_info": kline_message.topic_info,
                                "symbol": kline_message.symbol,
                                "id_info": kline_message.id_info,
                                "timestamp_or_version": kline_message.timestamp_or_version,
                            })
                        );
                        return Some(response_data);
                    } else {
                        info!("顶层 K 线消息结构解析成功，但未找到嵌套的 kline_data 字段 (Tag 308)");
                        // 这可能是一个只有顶层字段的控制消息
                        return Some(Response::new(
                            kline_message.topic_info.clone(), // 使用解析到的 Topic 信息
                            200,
                            "OK (Control Message)".to_string(),
                            json!({
                                "topic_info": kline_message.topic_info,
                                "symbol": kline_message.symbol,
                                "id_info": kline_message.id_info,
                                "timestamp_or_version": kline_message.timestamp_or_version,
                            })
                        ));
                    }
                }
                Err(e) => {
                    error!("尝试解析为 PublicSpotKlineV3ApiMessage 失败: {:?}", e);

                    return Some(Response::new("".to_string(), 500, format!("Protobuf K 线顶层消息解析出错: {:?}", e), Value::Null));
                }
            }
        }

        // 2. 尝试解析深度数据 (使用新的结构体)
        if prefix_string.contains("spot@public.aggre.depth.v3.api.pb") {
            // info!("通过 Topic 前缀判断为深度数据");

            // 尝试解析为 PublicIncreaseDepthsV3ApiMessage (新的顶层深度消息)
            match PublicIncreaseDepthsV3ApiMessage::decode(&po[..]) {
                Ok(depth_message) => {
                    // info!("成功解析为顶层深度消息结构");

                    // 检查是否包含嵌套的 depth_data 字段 (Tag 313)
                    if let Some(depth_data_content) = depth_message.depth_data {
                        // info!("找到并成功访问嵌套的 DepthDataContentV3");

                        // 填充 Response 并返回
                        let response_data = Response::new(
                            depth_message.topic_info.clone(), // 使用解析到的 Topic
                            200, "OK".to_string(),
                            serde_json::json!({
                                 // 嵌套消息内部的字段
                                "asks": depth_data_content.asks.into_iter().map(|item| serde_json::json!({"price": item.price, "quantity": item.quantity})).collect::<Vec<_>>(),
                                "bids": depth_data_content.bids.into_iter().map(|item| serde_json::json!({"price": item.price, "quantity": item.quantity})).collect::<Vec<_>>(),
                                "eventType": depth_data_content.event_type,
                                "version": depth_data_content.version,
                                "lastUpdateId": depth_data_content.last_update_id, // 新增字段

                                // 顶层字段
                                 "topic_info": depth_message.topic_info,
                                "symbol": depth_message.symbol,
                                 "timestamp": depth_message.timestamp, // 新增字段
                            })
                        );
                        return Some(response_data);

                    } else {
                        info!("顶层深度消息结构解析成功，但未找到嵌套的 depth_data 字段 (Tag 313)");
                        // 处理只有顶层字段的深度相关消息
                        return Some(Response::new(
                            depth_message.topic_info.clone(),
                            200, "OK (Control Message)".to_string(),
                            serde_json::json!({
                               "topic_info": depth_message.topic_info,
                               "symbol": depth_message.symbol,
                               "timestamp": depth_message.timestamp,
                           })
                        ));
                    }
                }
                Err(e) => {
                    error!("解析深度消息 PublicIncreaseDepthsV3ApiMessage 失败: {:?}", e);
                    // 保存数据以供分析
                    let file_path = Path::new("depth_error_data.bin");
                    // ... 保存 po 到文件 ...
                    return Some(Response::new("".to_string(), 500, format!("Protobuf 深度消息解析出错: {:?}", e), Value::Null));
                }
            }
        }

        // 如果都不是已知的 Protobuf 类型，处理未知消息
        error!("无法将二进制消息解析为任何已知 Protobuf 类型");
        // *** 在这里将原始二进制数据保存到文件 ***
        let file_path = Path::new("un_decode.bin");
        match File::create(&file_path) {
            Ok(mut file) => {
                match file.write_all(&po) {
                    Ok(_) => info!("原始 K 线二进制数据保存到 {:?}", file_path),
                    Err(write_e) => error!("保存 K 线二进制数据失败: {:?}", write_e),
                }
            }
            Err(create_e) => error!("创建文件 {:?} 失败: {:?}", file_path, create_e),
        }
        Some(Response::new("".to_string(), 400, "无法解析未知二进制消息".to_string(), Value::Null))
    }
    //数据解析
    pub fn ok_text(text: String) -> Response
    {
        info!("{}", text);
        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.data = json_value.clone();
                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()),
            MexcSpotWsSubscribeType::PuFuturesRecords("Min3".to_string()),
            MexcSpotWsSubscribeType::PuFuturesDepth
        ]);

        ws.set_symbols(vec!["BTC_USDT".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("链接失败");
    }
}