extended/src/main.rs

mod utils;
mod exchange;
mod strategy;
mod data_manager;

use anyhow::{Result};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use backtrace::Backtrace;
use tokio::spawn;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tracing::{error, info, warn};
use utils::log_setup;
use crate::data_manager::DataManager;
use crate::exchange::extended_account::ExtendedAccount;
use crate::exchange::extended_rest_client::ExtendedRestClient;
use crate::exchange::extended_stream_client::ExtendedStreamClient;
use crate::strategy::Strategy;
use crate::utils::response::Response;

#[tokio::main]
async fn main() {
    let _guards = log_setup::setup_logging().unwrap();

    // 主进程控制
    let running = Arc::new(AtomicBool::new(true));

    // panic错误捕获，panic级别的错误直接退出
    let panic_running = running.clone();
    std::panic::set_hook(Box::new(move |panic_info| {
        let msg = format!(
            "type=panic, msg={:?}, location={:?}",
            panic_info.to_string(),
            panic_info.location()
        );

        // 生成并格式化完整的堆栈跟踪
        let backtrace = Backtrace::new();
        let stack_trace = format!("{:?}", backtrace);

        // 一并打印堆栈跟踪
        warn!("{}\nStack Trace:\n{}", msg, stack_trace);
        panic_running.store(false, Ordering::Relaxed);
    }));

    // ---- 优雅停机处理 (示例: SIGINT/Ctrl+C) ----
    //注意：Windows上可能不支持所有信号，SIGINT通常可用
    let r = running.clone(); // 克隆 Arc 用于 SIGHUP/SIGTERM/SIGINT 处理
    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.expect("设置 Ctrl+C 处理器失败");
        warn!("接收到退出信号 (Ctrl+C)... 开始关闭.");
        r.store(false, Ordering::Relaxed);
    });

    // ---- 运行核心订阅逻辑 ----
    info!("==================================== 应用程序启动 =======================================");
    let task_running = running.clone();
    // 启动一个后台任务来执行订阅和数据处理
    let subscription_handle = tokio::spawn(async move {
        // 运行获取交易对和订阅 K 线的函数
        if let Err(e) = run_extended_subscriptions(task_running.clone()).await {
            error!("运行 Extended 订阅任务失败: {:?}", e);
            task_running.store(false, Ordering::Relaxed); // 如果启动失败，也设置停止标志
        }
    });
    info!("主循环开始，等待退出信号...");
    // ---- 主循环 ----
    // 保持主线程活动，等待 running 标志变为 false
    while running.load(Ordering::Relaxed) {
        // 可以添加一些周期性检查或任务，但主要是等待
        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    }

    info!("应用程序正在关闭...");

    // ---- 清理和关闭 ----
    // 等待订阅任务结束（如果它设计为可结束的话）
    info!("等待订阅任务完成...");
    // 可以给 subscription_handle 设置一个超时等待
    match tokio::time::timeout(tokio::time::Duration::from_secs(10), subscription_handle).await {
        Ok(Ok(_)) => info!("订阅任务正常结束。"),
        Ok(Err(e)) => error!("订阅任务返回错误: {:?}", e),
        Err(_) => warn!("等待订阅任务超时。"),
    }

    info!("应用程序已关闭。");
}

/// 运行 Extended 的主要订阅任务
///
/// # Arguments
/// * `running` - 用于控制程序是否继续运行的原子布尔值 (Arc 包裹)
///
/// # Returns
pub async fn run_extended_subscriptions(running: Arc<AtomicBool>) -> Result<()> {
    let market = "BTC-USD";
    let account = ExtendedAccount::new(
        "9ae4030902ab469a1bae8a90464e2e91",
        "0x71e16e49b717b851ced8347cf0dfa8f490bfb826323b9af624a66285dc99672",
        "0x47cdde8952945c13460f9129644eade096100810fba59de05452b34aacecff6",
        220844,
    );
    let is_testnet = false;

    // 订阅数据的客户端
    let stream_client_list = vec![
        ExtendedStreamClient::best_prices(format!("ExtendedBestPrices_{}", market).as_str(), None, market, is_testnet),
        ExtendedStreamClient::account("ExtendedAccount", Some(account.clone()), is_testnet),
    ];
    
    // rest客户端
    let rest_client = ExtendedRestClient::new("ExtendedRestClient", Some(account), market, is_testnet).await?;
    let rest_client_am = Arc::new(Mutex::new(rest_client));

    // 数据管理及消息分发
    let data_manager = DataManager::new();
    let data_manager_am = Arc::new(Mutex::new(data_manager));

    // 策略执行
    let strategy = Strategy::new(rest_client_am.clone());
    let strategy_am = Arc::new(Mutex::new(strategy));

    // 异步去订阅、并阻塞
    for mut stream_client in stream_client_list {
        let running_clone = Arc::clone(&running);

        // 定义需要处理数据的fun
        let dm = data_manager_am.clone();
        let sm = strategy_am.clone();
        let fun = move |response: Response| {
            if response.code != 200 {
                error!("出现错误代码：{}", serde_json::to_string_pretty(&response.data).unwrap());
            }

            let dm_clone = Arc::clone(&dm);
            let sm_clone = Arc::clone(&sm);
            async move {
                // 数据不新鲜直接跳过
                if response.reach_time - response.received_time > 100 {
                    return
                }
                
                let mut dm_guard = dm_clone.lock().await;

                // 记录消息延迟
                dm_guard.record_latency(response.received_time, response.reach_time);

                // 交给消息分发函数，并在此处消费掉错误消息
                if let Err(e) = dm_guard.dispatch_message(&response).await {
                    warn!("消息分发过程中出现错误: {}", e);
                }

                // 随后执行策略
                let mut sm_guard = sm_clone.lock().await;
                if let Err(e) = sm_guard.do_strategy(&dm_guard).await {
                    warn!("策略执行过程中出现错误: {}", e);
                }
            }
        };

        // 这个通道主要是为了后面给这个ws发送消息
        let (write_tx, write_rx) = futures_channel::mpsc::unbounded::<Message>();
        let write_tx_am = Arc::new(Mutex::new(write_tx));

        spawn(async move {
            // 链接
            stream_client.ws_connect_async(running_clone, fun, &write_tx_am, write_rx)
                .await
                .expect("ws链接失败");
        });
    }

    // // 网络延迟统计
    // let running_clone = Arc::clone(&running);
    // spawn(async move {
    //     let mut interval = tokio::time::interval(Duration::from_secs(10));
    //
    //     while running_clone.load(Ordering::SeqCst) {
    //         interval.tick().await; // 等待下一个周期
    //
    //         if !running_clone.load(Ordering::SeqCst) {
    //             break;
    //         }
    //
    //         let mut total_delay_sum = 0i64;
    //         let mut total_message_count = 0u64;
    //
    //         // --- 第一步：收集 DataManager 的统计数据 ---
    //         // 直接调用 DataManager 的方法
    //         let manager_lock = data_manager_am.lock().await; // 锁定单个 DataManager
    //         let (current_sum, current_count) = manager_lock.get_delay_stats(); // 使用原子读
    //
    //         // 使用 saturating_add 防止聚合时溢出 (虽然 u64 很大，但好习惯)
    //         total_delay_sum = total_delay_sum.saturating_add(current_sum);
    //         total_message_count = total_message_count.saturating_add(current_count);
    //
    //         manager_lock.reset_delay_stats(); // 调用重置方法 (使用原子写)
    //
    //         // --- 第二步：计算并报告平均延迟 ---
    //         if total_message_count > 0 {
    //             let average_delay = total_delay_sum as f64 / total_message_count as f64;
    //             info!(
    //                     "当前 WS 平均延迟: {:.2} 毫秒 (基于 {} 条消息测量)",
    //                     average_delay, total_message_count
    //                 );
    //         } else {
    //             info!("WS 延迟报告：本周期内未收到用于计算延迟的消息。");
    //         }
    //     }
    //     info!("延迟报告任务已停止。");
    // });

    Ok(())
}