extended/src/exchange/extended_stream_client.rs

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

use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use serde_json::json;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{http, Message};
use tracing::{error, info, trace, warn};
use anyhow::Result;
use chrono::Utc;
use tokio_tungstenite::tungstenite::handshake::client::{generate_key, Request};
use tracing_subscriber::fmt::format::json;
use crate::exchange::extended_account::ExtendedAccount;
use crate::utils::response::Response;
use crate::utils::stream_utils::{StreamUtils, HeartbeatType};

#[derive(Clone)]
#[allow(dead_code)]
pub struct ExtendedStreamClient {
    // 标签
    tag: String,
    // 地址
    address_url: String,
    // 账号
    account_option: Option<ExtendedAccount>,
    // 心跳间隔
    heartbeat_time: u64,
}

impl ExtendedStreamClient {
    // ============================================= 构造函数 ================================================
    fn new(tag: String, account_option: Option<ExtendedAccount>, subscribe_pattern: String) -> ExtendedStreamClient {
        let host = "wss://api.starknet.extended.exchange/stream.extended.exchange/v1/".to_string();  // mainnet
        // let host = "wss://api.starknet.sepolia.extended.exchange/stream.extended.exchange/v1/".to_string();  // testnet
        
        let address_url = format!("{}{}", host, subscribe_pattern);

        ExtendedStreamClient {
            tag,
            address_url,
            account_option,
            heartbeat_time: 1000 * 10,
        }
    }

    // ============================================= 订阅函数 ================================================
    pub fn order_books(tag: String, account_option: Option<ExtendedAccount>, symbol: String) -> ExtendedStreamClient {
        Self::new(tag, account_option, format!("orderbooks/{}", symbol))
    }

    // 链接
    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<()>
        where
            F: Fn(Response) -> Future + Clone + Send + 'static + Sync,
            Future: std::future::Future<Output=()> + Send + 'static, // 确保 Fut 是一个 Future，且输出类型为 ()
    {
        let address_url = self.address_url.clone();
        let tag = self.tag.clone();

        // 自动心跳包
        let write_tx_clone1 = write_tx_am.clone();
        let heartbeat_time = self.heartbeat_time.clone();
        tokio::spawn(async move {
            let ping_obj = json!({"method":"PING"});
            StreamUtils::ping_pong(write_tx_clone1, HeartbeatType::Custom(ping_obj.to_string()), heartbeat_time).await;
        });

        if self.account_option.is_some() {
            // 登录相关
        }

        // 提取host
        let parsed_uri: http::Uri = address_url.parse()?;
        let host_domain = parsed_uri.host().ok_or("URI 缺少主机名").unwrap().to_string();
        let host_header_value = if let Some(port) = parsed_uri.port_u16() {
            // 如果端口不是默认的 80 (for ws) 或 443 (for wss)，则需要包含端口
            // 这里只是简单地判断，更严谨的判断可以根据 scheme 来
            match parsed_uri.scheme_str() {
                Some("ws") if port == 80 => host_domain.to_string(),
                Some("wss") if port == 443 => host_domain.to_string(),
                _ => format!("{}:{}", host_domain, port), // 否则包含端口
            }
        } else {
            host_domain.to_string() // 没有端口或使用默认端口
        };

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

            loop {
                // 通过构建request的方式进行ws链接，可以携带header
                let request = Request::builder()
                    .method("GET")
                    .uri(&address_url)
                    .header("Sec-WebSocket-Key", generate_key())
                    .header("Sec-WebSocket-Version", "13")
                    .header("Host", host_header_value.clone())
                    .header("User-Agent", "RustClient/1.0")
                    .header("Upgrade", "websocket")
                    .header("Connection", "Upgrade")
                    .body(())
                    .unwrap();

                trace!("Extended_usdt_swap socket 连接中……");
                StreamUtils::ws_connect_async(is_shutdown_arc.clone(), handle_function.clone(), request,
                                              false, tag.clone(), vec![], write_to_socket_rx_arc.clone(),
                                              Self::message_text, Self::message_ping, Self::message_pong, Self::message_binary).await;

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

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

        Ok(())
    }

    //数据解析-Text
    pub fn message_text(text: String) -> Option<Response> {
        let mut res_data = Response::new("".to_string(), -201, "success".to_string(), Value::Null);
        let json_value: Value = serde_json::from_str(&text).unwrap();
        
        // info!("等待解析：{}", serde_json::to_string_pretty(&json_value).unwrap());

        match json_value["ts"].as_i64() {
            Some(ts) => {
                res_data.reach_time = ts;
                res_data.received_time = Utc::now().timestamp_millis();
                res_data.code = 200;
                res_data.data = json_value.clone();
            }
            None => {
                res_data.data = json_value.clone();
                res_data.code = -1;
                res_data.message = text;
            }
        }

        Option::from(res_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,
        //                     "success".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": kline_message.timestamp,
        //                         })
        //                 );
        //                 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": kline_message.timestamp,
        //                     })
        //                 ));
        //             }
        //         }
        //         Err(e) => {
        //             error!("尝试解析为 PublicSpotKlineV3ApiMessage 失败: {:?}", e);
        //         }
        //     }
        // }
        //
        // // 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,
        //                     "success".to_string(),
        //                     json!({
        //                          // 嵌套消息内部的字段
        //                         "asks": depth_data_content.asks.into_iter().map(|item| json!({"price": item.price, "quantity": item.quantity})).collect::<Vec<_>>(),
        //                         "bids": depth_data_content.bids.into_iter().map(|item| 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);
        //         }
        //     }
        // }

        // 如果都不是已知的 Protobuf 类型，处理未知消息
        error!("无法将二进制消息解析为任何已知 Protobuf 类型, {}", prefix_string);

        Some(Response::new("".to_string(), 400, "无法解析未知二进制消息".to_string(), Value::Null))
    }
}

#[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::extended_stream_client::{ExtendedStreamClient};
    use crate::utils::response::Response;
    use crate::utils::log_setup::setup_logging;

    #[tokio::test]
    async fn test_extended_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 = ExtendedStreamClient::order_books("Extended".to_string(), None, "BTC-USD".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("链接失败");
    }
}