import axios from 'axios';
import WebSocket from 'ws';

// 配置
const CONFIG = {
    baseURL: 'http://localhost:8888',          // API 基础地址
    wsURL: 'ws://localhost:8888/clob/ws',      // WebSocket 地址
    marketId: 'market_123',                    // 默认市场 ID (如果没有从 API 获取到)
    tokenId: 'market_123',                     // 默认 Token ID (如果没有从 API 获取到)
    userCount: 20,                             // 并发用户数量
    ordersPerUser: 100,                        // 每个用户的下单总数
    orderIntervalMs: 1000,                     // 下单间隔 (毫秒)
    basePrice: 5000,                           // 基准价格 (例如: 0.50，如果精度是 10000)
    priceSpread: 10,                           // 价格波动范围 (+/- 10 个步长，每个步长 100)
    baseSize: 1000000,                         // 基准数量 (1 个单位)
    sizeSpread: 5                              // 数量波动范围 (+/- 5 个单位)
};

// 市场响应接口定义
interface PmMarket {
    ID: number;
    clobTokenIds: string[]; // 数据库中可能是 JSON 字符串，但 API 可能返回解析后的对象或数组
}

// 用户会话状态接口
interface UserSession {
    index: number;              // 用户索引 (1-based)
    walletAddress: string;      // 钱包地址
    token: string;              // 登录 Token
    userId: number;             // 用户 ID
    ordersPlaced: number;       // 已下单数量
    successCount: number;       // 下单成功数量
    failCount: number;          // 下单失败数量
    wsMessages: number;         // 接收到的 WebSocket 消息数量
    totalLatency: number;       // 总延迟 (毫秒，用于计算平均值)
    intervalId?: NodeJS.Timeout;// 定时器 ID
    ws?: WebSocket;             // WebSocket 连接实例
}

// 生成确定性的测试钱包地址
// 格式: 0x0000... + 索引 (补全到 40 个字符)
// 这种格式用于配合后端的测试模式，绕过签名验证
function getTestWalletAddress(index: number): string {
    const hexIndex = index.toString(16).padStart(36, '0');
    return `0x0000${hexIndex}`;
}

// 用户登录
// 使用测试钱包地址登录，获取 Token
async function loginUser(index: number): Promise<UserSession | null> {
    const walletAddress = getTestWalletAddress(index + 1); // 从 1 开始编号
    
    try {
        // 准备登录负载
        // 服务器逻辑: 如果是测试环境且钱包地址以 "0x0000" 开头，则跳过签名验证
        const payload = {
            walletAddress: walletAddress,
            nonce: "test_nonce",
            signature: "test_signature",
            message: "test_message"
        };

        const response = await axios.post(`${CONFIG.baseURL}/base/walletLogin`, payload);
        
        if (response.data.code === 0 && response.data.data) {
            const userData = response.data.data.user;
            const token = response.data.data.token;
            // console.log(`[用户 ${index + 1}] 登录成功。ID: ${userData.ID}`);
            
            return {
                index: index + 1,
                walletAddress,
                token,
                userId: userData.ID,
                ordersPlaced: 0,
                successCount: 0,
                failCount: 0,
                wsMessages: 0,
                totalLatency: 0
            };
        } else {
            console.error(`[用户 ${index + 1}] 登录失败:`, response.data.msg);
            return null;
        }
    } catch (error: any) {
        console.error(`[用户 ${index + 1}] 登录错误:`, error.message);
        return null;
    }
}

// 建立 WebSocket 连接并订阅市场
function connectWs(session: UserSession, tokenId: string) {
    // 当前系统 WebSocket 连接不需要鉴权 (基于 public group)
    
    const ws = new WebSocket(CONFIG.wsURL);
    
    ws.on('open', () => {
        // 连接打开后发送订阅消息
        const msg = {
            type: "market",
            assets_ids: [tokenId]
        };
        ws.send(JSON.stringify(msg));
    });

    ws.on('message', (data: WebSocket.Data) => {
        session.wsMessages++;
        // const msg = JSON.parse(data.toString());
        // console.log(`[用户 ${session.index}] WS 消息:`, msg);
    });

    ws.on('error', (err) => {
        // console.error(`[用户 ${session.index}] WS 错误:`, err.message);
    });

    session.ws = ws;
}

// 执行下单操作
async function placeOrder(session: UserSession) {
    // 检查是否达到下单上限
    if (session.ordersPlaced >= CONFIG.ordersPerUser) {
        if (session.intervalId) {
            clearInterval(session.intervalId);
        }
        if (session.ws) {
            session.ws.close();
        }
        return;
    }

    // 方向: 随机买入/卖出 (1=买, 2=卖)
    const isBuy = Math.random() > 0.5;
    const side = isBuy ? 1 : 2; 
    
    // 价格: 基准价格 +/- 随机波动 (步长 100)
    // 5000 +/- (0..10)*100
    const priceOffset = Math.floor(Math.random() * CONFIG.priceSpread) * 100;
    const price = isBuy ? CONFIG.basePrice - priceOffset : CONFIG.basePrice + priceOffset;
    
    // 数量: 基准数量 +/- 随机波动 (步长 1000000)
    const sizeOffset = Math.floor(Math.random() * CONFIG.sizeSpread) * 1000000;
    const size = CONFIG.baseSize + sizeOffset;

    const payload = {
        tokenID: CONFIG.tokenId,
        side: side,
        price: price,
        size: size,
        feeRateBps: 10,
        nonce: Date.now(), // 简单的 nonce
        expiration: 0,     // 0 表示不过期 (GTC)
        orderType: 0       // 0: GTC (Good Till Cancel)
    };

    const start = Date.now();
    try {
        const response = await axios.post(`${CONFIG.baseURL}/clob/gateway/submitOrder`, payload, {
            headers: {
                'x-token': session.token
            }
        });
        
        const latency = Date.now() - start;
        session.totalLatency += latency;

        session.ordersPlaced++;

        if (response.data.code === 0) {
            session.successCount++;
        } else {
            session.failCount++;
            // 仅打印前 5 个错误，避免日志刷屏
            if (session.failCount <= 5) {
                console.error(`[用户 ${session.index}] 下单失败:`, response.data.msg);
            }
        }

    } catch (error: any) {
        session.ordersPlaced++;
        session.failCount++;
        // 仅打印前 5 个错误
        if (session.failCount <= 5) {
            let msg = error.message;
            if (error.response) {
                msg += ` 状态码: ${error.response.status}`;
            }
            console.error(`[用户 ${session.index}] 下单错误:`, msg);
        }
    }
}

// 主函数：运行负载测试
async function runLoadTest() {
    console.log(`开始负载测试，并发用户数: ${CONFIG.userCount}...`);
    console.log(`目标: 每用户 ${CONFIG.ordersPerUser} 单`);
    console.log(`速率: 每 ${CONFIG.orderIntervalMs}ms 一单`);
    
    // 1. 所有用户登录
    const sessions: UserSession[] = [];
    const loginPromises = [];
    
    for (let i = 0; i < CONFIG.userCount; i++) {
        loginPromises.push(loginUser(i).then(session => {
            if (session) sessions.push(session);
        }));
        // 稍微错开登录时间，避免瞬间压垮数据库连接池
        await new Promise(r => setTimeout(r, 50));
    }
    
    await Promise.all(loginPromises);

    console.log(`\n成功登录 ${sessions.length}/${CONFIG.userCount} 个用户。`);

    // 2. 获取市场列表以找到有效的 Token ID
    if (sessions.length > 0) {
        try {
            console.log("正在获取市场列表以查找有效 Token ID...");
            const marketRes = await axios.get(`${CONFIG.baseURL}/PmMarket/getPmMarketList?page=1&pageSize=1`, {
                headers: {
                    'x-token': sessions[0].token
                }
            });

            if (marketRes.data.code === 0 && marketRes.data.data.list && marketRes.data.data.list.length > 0) {
                const market = marketRes.data.data.list[0];
                console.log(`找到市场: ${market.ID} - ${market.question}`);
                
                // 解析 clobTokenIds
                let tokenIds = market.clobTokenIds;
                if (typeof tokenIds === 'string') {
                    try {
                        tokenIds = JSON.parse(tokenIds);
                    } catch (e) {
                        // console.error("解析 clobTokenIds JSON 失败:", e);
                    }
                }

                if (Array.isArray(tokenIds) && tokenIds.length > 0) {
                    CONFIG.tokenId = tokenIds[0];
                    CONFIG.marketId = market.ID.toString();
                    console.log(`使用配置 -> MarketID: ${CONFIG.marketId}, TokenID: ${CONFIG.tokenId}`);
                } else {
                    console.error("市场没有有效的 clobTokenIds 数组:", tokenIds);
                }
            } else {
                console.error("未找到市场或 API 错误:", marketRes.data);
            }
        } catch (error: any) {
            console.error("获取市场列表失败:", error.message);
        }
    }

    // 3. 所有用户建立 WebSocket 连接
    console.log("正在连接 WebSocket...");
    sessions.forEach(s => connectWs(s, CONFIG.tokenId));

    console.log('开始下单循环...\n');

    const startTime = Date.now();
    
    // 4. 启动下单循环
    // 使用 Promise 跟踪每个用户的任务完成情况
    const userTasks = sessions.map(session => {
        return new Promise<void>((resolve) => {
            // 初始随机延迟，错开启动时间
            setTimeout(() => {
                session.intervalId = setInterval(async () => {
                    if (session.ordersPlaced >= CONFIG.ordersPerUser) {
                        if (session.intervalId) clearInterval(session.intervalId);
                        if (session.ws) session.ws.close();
                        resolve();
                        return;
                    }
                    await placeOrder(session);
                }, CONFIG.orderIntervalMs);
            }, Math.random() * 1000); // 1秒内的随机延迟
        });
    });

    // 5. 状态监控 (每秒刷新)
    const monitorInterval = setInterval(() => {
        const totalPlaced = sessions.reduce((sum, s) => sum + s.ordersPlaced, 0);
        const totalSuccess = sessions.reduce((sum, s) => sum + s.successCount, 0);
        const totalFail = sessions.reduce((sum, s) => sum + s.failCount, 0);
        const totalWs = sessions.reduce((sum, s) => sum + s.wsMessages, 0);
        const totalLatency = sessions.reduce((sum, s) => sum + s.totalLatency, 0);
        const avgLatency = totalSuccess > 0 ? (totalLatency / totalSuccess).toFixed(0) : "0";

        const elapsed = (Date.now() - startTime) / 1000;
        const rate = elapsed > 0 ? (totalPlaced / elapsed).toFixed(1) : "0.0";
        
        // 覆盖当前行输出状态
        process.stdout.write(`\r[状态] 时间: ${elapsed.toFixed(1)}s | 订单: ${totalPlaced}/${CONFIG.userCount * CONFIG.ordersPerUser} | 速率: ${rate} ops | 成功: ${totalSuccess} | 失败: ${totalFail} | WS消息: ${totalWs} | 延迟: ${avgLatency}ms`);
        
        if (totalPlaced >= CONFIG.userCount * CONFIG.ordersPerUser) {
            clearInterval(monitorInterval);
            process.stdout.write('\n'); // 完成后换行
        }
    }, 1000);

    // 等待所有用户完成
    await Promise.all(userTasks);
    
    const endTime = Date.now();
    const duration = (endTime - startTime) / 1000;
    
    console.log('\n\n负载测试完成!');
    console.log(`总耗时: ${duration.toFixed(2)}s`);
    
    const totalSuccess = sessions.reduce((sum, s) => sum + s.successCount, 0);
    const totalFail = sessions.reduce((sum, s) => sum + s.failCount, 0);
    const totalWs = sessions.reduce((sum, s) => sum + s.wsMessages, 0);
    console.log(`总订单数: ${totalSuccess + totalFail} (成功: ${totalSuccess}, 失败: ${totalFail})`);
    console.log(`总共接收 WS 消息数: ${totalWs}`);
}

// 运行测试
runLoadTest().catch(console.error);