// Sub2API Gateway 负载测试脚本
// 使用 k6 进行性能测试
// 运行: k6 run --env BASE_URL=http://localhost:8080 --env API_KEY=your_key gateway-load-test.js

import http from 'k6/http';
import { check, sleep, group } from 'k6';
import { Rate, Trend, Counter, Gauge } from 'k6/metrics';
import { randomIntBetween } from 'https://jslib.k6.io/k6-utils/1.2.0/index.js';

// ============================================
// 自定义指标定义
// ============================================
const errorRate = new Rate('error_rate');
const successRate = new Rate('success_rate');
const responseTime = new Trend('response_time');
const responseTimeP95 = new Trend('response_time_p95');
const responseTimeP99 = new Trend('response_time_p99');
const requestCounter = new Counter('total_requests');
const successCounter = new Counter('success_requests');
const failureCounter = new Counter('failure_requests');
const activeVUs = new Gauge('active_vus');

// ============================================
// 测试配置
// ============================================
export const options = {
  // 测试阶段配置
  stages: [
    { duration: '2m', target: 50 },    // 渐进加压: 0 -> 50 VU
    { duration: '5m', target: 50 },    // 稳定期: 保持50 VU
    { duration: '2m', target: 100 },   // 峰值测试: 50 -> 100 VU
    { duration: '5m', target: 100 },   // 峰值稳定: 保持100 VU
    { duration: '2m', target: 200 },   // 压力测试: 100 -> 200 VU
    { duration: '3m', target: 200 },   // 压力稳定: 保持200 VU
    { duration: '2m', target: 0 },     // 减压: 200 -> 0 VU
  ],

  // 性能阈值配置
  thresholds: {
    // HTTP请求持续时间
    http_req_duration: ['p(50)<100', 'p(95)<500', 'p(99)<2000'],
    http_req_failed: ['rate<0.01'],           // 错误率 < 1%
    error_rate: ['rate<0.001'],                // 自定义错误率 < 0.1%
    success_rate: ['rate>0.99'],               // 成功率 > 99%
    
    // 特定端点的阈值
    'response_time{endpoint:models}': ['p(95)<200'],
    'response_time{endpoint:chat_completions}': ['p(95)<5000'],
    'response_time{endpoint:messages}': ['p(95)<5000'],
  },

  // 其他配置
  discardResponseBodies: true,  // 丢弃响应体以节省内存
  maxRedirects: 4,
  userAgent: 'Sub2API-LoadTest/1.0',
};

// ============================================
// 环境变量配置
// ============================================
const BASE_URL = __ENV.BASE_URL || 'http://localhost:8080';
const API_KEY = __ENV.API_KEY || 'sk-test-key';
const ENABLE_STREAMING = __ENV.ENABLE_STREAMING === 'true';

// ============================================
// 测试数据
// ============================================
const TEST_MODELS = [
  'claude-3-opus-20240229',
  'claude-3-sonnet-20240229',
  'claude-3-haiku-20240307',
  'gpt-4-turbo-preview',
  'gpt-4',
  'gpt-3.5-turbo',
];

const TEST_MESSAGES = [
  [{ role: 'user', content: 'Hello, how are you?' }],
  [{ role: 'user', content: 'What is the weather like today?' }],
  [{ role: 'user', content: 'Explain quantum computing in simple terms.' }],
  [{ role: 'user', content: 'Write a short poem about technology.' }],
  [{ role: 'system', content: 'You are a helpful assistant.' }, { role: 'user', content: 'Help me with my homework.' }],
];

// ============================================
// 设置函数 - 每个VU执行一次
// ============================================
export function setup() {
  console.log(`Starting load test against ${BASE_URL}`);
  console.log(`API Key: ${API_KEY.substring(0, 10)}...`);
  
  // 验证API Key有效性
  const res = http.get(`${BASE_URL}/v1/models`, {
    headers: {
      'Authorization': `Bearer ${API_KEY}`,
      'Content-Type': 'application/json',
    },
  });
  
  if (res.status !== 200) {
    console.error(`API Key validation failed: ${res.status}`);
    return { valid: false };
  }
  
  console.log('API Key validated successfully');
  return { valid: true, startTime: new Date().toISOString() };
}

// ============================================
// 主测试函数 - 每个VU循环执行
// ============================================
export default function (data) {
  if (!data.valid) {
    console.error('Invalid setup, skipping test');
    return;
  }

  activeVUs.add(__VU);

  // 随机选择测试场景 (权重分配)
  const scenario = randomIntBetween(1, 100);
  
  if (scenario <= 40) {
    // 40% - 模型列表查询 (轻量级)
    testListModels();
  } else if (scenario <= 70) {
    // 30% - Chat Completions
    testChatCompletions();
  } else if (scenario <= 90) {
    // 20% - Messages (Claude API)
    testMessages();
  } else {
    // 10% - 混合操作
    testMixedOperations();
  }

  // 随机休眠 100ms - 1000ms 模拟真实用户行为
  sleep(randomIntBetween(1, 10) / 10);
}

// ============================================
// 测试场景: 获取模型列表
// ============================================
function testListModels() {
  group('List Models', () => {
    const startTime = new Date();
    
    const res = http.get(`${BASE_URL}/v1/models`, {
      headers: {
        'Authorization': `Bearer ${API_KEY}`,
        'Content-Type': 'application/json',
      },
      tags: { endpoint: 'models' },
    });

    const duration = new Date() - startTime;
    responseTime.add(duration, { endpoint: 'models' });
    requestCounter.add(1, { endpoint: 'models' });

    const success = check(res, {
      'models_status_is_200': (r) => r.status === 200,
      'models_response_has_data': (r) => {
        try {
          const body = JSON.parse(r.body);
          return body.data && Array.isArray(body.data);
        } catch (e) {
          return false;
        }
      },
      'models_response_time_ok': () => duration < 500,
    });

    if (success) {
      successCounter.add(1, { endpoint: 'models' });
      successRate.add(true);
      errorRate.add(false);
    } else {
      failureCounter.add(1, { endpoint: 'models' });
      successRate.add(false);
      errorRate.add(true);
      console.error(`Models request failed: ${res.status}, body: ${res.body}`);
    }
  });
}

// ============================================
// 测试场景: Chat Completions
// ============================================
function testChatCompletions() {
  group('Chat Completions', () => {
    const model = TEST_MODELS[randomIntBetween(0, TEST_MODELS.length - 1)];
    const messages = TEST_MESSAGES[randomIntBetween(0, TEST_MESSAGES.length - 1)];
    
    const payload = JSON.stringify({
      model: model,
      messages: messages,
      max_tokens: randomIntBetween(50, 200),
      temperature: 0.7,
      stream: ENABLE_STREAMING,
    });

    const startTime = new Date();
    
    const res = http.post(`${BASE_URL}/v1/chat/completions`, payload, {
      headers: {
        'Authorization': `Bearer ${API_KEY}`,
        'Content-Type': 'application/json',
      },
      tags: { endpoint: 'chat_completions' },
      timeout: '30s',
    });

    const duration = new Date() - startTime;
    responseTime.add(duration, { endpoint: 'chat_completions' });
    requestCounter.add(1, { endpoint: 'chat_completions' });

    const success = check(res, {
      'chat_status_is_200': (r) => r.status === 200,
      'chat_response_valid': (r) => {
        try {
          const body = JSON.parse(r.body);
          return body.choices && body.choices.length > 0;
        } catch (e) {
          return false;
        }
      },
      'chat_response_time_acceptable': () => duration < 10000, // 10s
    });

    if (success) {
      successCounter.add(1, { endpoint: 'chat_completions' });
      successRate.add(true);
      errorRate.add(false);
    } else {
      failureCounter.add(1, { endpoint: 'chat_completions' });
      successRate.add(false);
      errorRate.add(true);
      console.error(`Chat completions failed: ${res.status}, duration: ${duration}ms`);
    }
  });
}

// ============================================
// 测试场景: Messages (Claude API)
// ============================================
function testMessages() {
  group('Messages', () => {
    const model = TEST_MODELS[randomIntBetween(0, 2)]; // Claude models only
    const messages = TEST_MESSAGES[randomIntBetween(0, TEST_MESSAGES.length - 1)];
    
    const payload = JSON.stringify({
      model: model,
      messages: messages,
      max_tokens: randomIntBetween(50, 200),
    });

    const startTime = new Date();
    
    const res = http.post(`${BASE_URL}/v1/messages`, payload, {
      headers: {
        'Authorization': `Bearer ${API_KEY}`,
        'Content-Type': 'application/json',
        'anthropic-version': '2023-06-01',
      },
      tags: { endpoint: 'messages' },
      timeout: '30s',
    });

    const duration = new Date() - startTime;
    responseTime.add(duration, { endpoint: 'messages' });
    requestCounter.add(1, { endpoint: 'messages' });

    const success = check(res, {
      'messages_status_is_200': (r) => r.status === 200,
      'messages_response_valid': (r) => {
        try {
          const body = JSON.parse(r.body);
          return body.content && body.content.length > 0;
        } catch (e) {
          return false;
        }
      },
      'messages_response_time_acceptable': () => duration < 10000,
    });

    if (success) {
      successCounter.add(1, { endpoint: 'messages' });
      successRate.add(true);
      errorRate.add(false);
    } else {
      failureCounter.add(1, { endpoint: 'messages' });
      successRate.add(false);
      errorRate.add(true);
      console.error(`Messages request failed: ${res.status}`);
    }
  });
}

// ============================================
// 测试场景: 混合操作
// ============================================
function testMixedOperations() {
  group('Mixed Operations', () => {
    // 1. 获取模型列表
    testListModels();
    sleep(0.5);
    
    // 2. 随机选择一个API调用
    if (randomIntBetween(1, 2) === 1) {
      testChatCompletions();
    } else {
      testMessages();
    }
  });
}

// ============================================
// 拆卸函数 - 测试结束时执行
// ============================================
export function teardown(data) {
  console.log('Load test completed');
  console.log(`Start time: ${data.startTime}`);
  console.log(`End time: ${new Date().toISOString()}`);
}

// ============================================
// 性能测试报告处理器
// ============================================
export function handleSummary(data) {
  return {
    'stdout': textSummary(data, { indent: ' ', enableColors: true }),
    'summary.json': JSON.stringify({
      metrics: {
        http_req_duration: data.metrics.http_req_duration,
        http_req_failed: data.metrics.http_req_failed,
        error_rate: data.metrics.error_rate,
        success_rate: data.metrics.success_rate,
      },
      thresholds: data.thresholds,
      state: data.state,
    }, null, 2),
    'summary.html': generateHTMLReport(data),
  };
}

// 生成文本摘要
function textSummary(data, options) {
  const indent = options.indent || '';
  const colors = options.enableColors !== false;
  
  const red = colors ? '\x1b[31m' : '';
  const green = colors ? '\x1b[32m' : '';
  const yellow = colors ? '\x1b[33m' : '';
  const reset = colors ? '\x1b[0m' : '';
  
  let summary = '';
  summary += `${indent}╔════════════════════════════════════════════════════════╗\n`;
  summary += `${indent}║           Sub2API 负载测试报告                          ║\n`;
  summary += `${indent}╚════════════════════════════════════════════════════════╝\n\n`;
  
  summary += `${indent}测试持续时间: ${(data.state.testRunDurationMs / 1000).toFixed(2)}s\n`;
  summary += `${indent}虚拟用户数: ${data.metrics.vus ? data.metrics.vus.max : 'N/A'}\n`;
  summary += `${indent}总请求数: ${data.metrics.http_reqs ? data.metrics.http_reqs.count : 'N/A'}\n\n`;
  
  // 响应时间统计
  if (data.metrics.http_req_duration) {
    const dur = data.metrics.http_req_duration;
    summary += `${indent}响应时间统计:\n`;
    summary += `${indent}  平均: ${dur.avg.toFixed(2)}ms\n`;
    summary += `${indent}  最小: ${dur.min.toFixed(2)}ms\n`;
    summary += `${indent}  最大: ${dur.max.toFixed(2)}ms\n`;
    summary += `${indent}  P50: ${dur.med.toFixed(2)}ms\n`;
    summary += `${indent}  P95: ${dur['p(95)'].toFixed(2)}ms\n`;
    summary += `${indent}  P99: ${dur['p(99)'].toFixed(2)}ms\n\n`;
  }
  
  // 错误率
  if (data.metrics.http_req_failed) {
    const failRate = data.metrics.http_req_failed.rate * 100;
    const color = failRate > 1 ? red : (failRate > 0.1 ? yellow : green);
    summary += `${indent}错误率: ${color}${failRate.toFixed(4)}%${reset}\n`;
  }
  
  // 阈值检查结果
  summary += `\n${indent}阈值检查:\n`;
  for (const [name, result] of Object.entries(data.thresholds || {})) {
    const color = result.ok ? green : red;
    summary += `${indent}  ${color}${result.ok ? '✓' : '✗'} ${name}${reset}\n`;
  }
  
  return summary;
}

// 生成HTML报告
function generateHTMLReport(data) {
  return `
<!DOCTYPE html>
<html>
<head>
  <title>Sub2API 负载测试报告</title>
  <style>
    body { font-family: Arial, sans-serif; margin: 40px; background: #f5f5f5; }
    .container { max-width: 1200px; margin: 0 auto; background: white; padding: 30px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
    h1 { color: #333; border-bottom: 3px solid #4CAF50; padding-bottom: 10px; }
    .metric { display: inline-block; margin: 10px 20px 10px 0; padding: 15px; background: #f9f9f9; border-radius: 4px; }
    .metric-label { font-size: 12px; color: #666; text-transform: uppercase; }
    .metric-value { font-size: 24px; font-weight: bold; color: #333; }
    .success { color: #4CAF50; }
    .warning { color: #FF9800; }
    .error { color: #f44336; }
    table { width: 100%; border-collapse: collapse; margin: 20px 0; }
    th, td { padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }
    th { background: #4CAF50; color: white; }
    tr:hover { background: #f5f5f5; }
  </style>
</head>
<body>
  <div class="container">
    <h1>Sub2API 负载测试报告</h1>
    <p>生成时间: ${new Date().toLocaleString()}</p>
    
    <div class="metrics">
      <div class="metric">
        <div class="metric-label">总请求数</div>
        <div class="metric-value">${data.metrics.http_reqs ? data.metrics.http_reqs.count.toLocaleString() : 'N/A'}</div>
      </div>
      <div class="metric">
        <div class="metric-label">平均响应时间</div>
        <div class="metric-value">${data.metrics.http_req_duration ? data.metrics.http_req_duration.avg.toFixed(2) : 'N/A'}ms</div>
      </div>
      <div class="metric">
        <div class="metric-label">P95 响应时间</div>
        <div class="metric-value">${data.metrics.http_req_duration ? data.metrics.http_req_duration['p(95)'].toFixed(2) : 'N/A'}ms</div>
      </div>
      <div class="metric">
        <div class="metric-label">错误率</div>
        <div class="metric-value ${data.metrics.http_req_failed && data.metrics.http_req_failed.rate > 0.01 ? 'error' : 'success'}">
          ${data.metrics.http_req_failed ? (data.metrics.http_req_failed.rate * 100).toFixed(4) : 'N/A'}%
        </div>
      </div>
    </div>
    
    <h2>阈值检查结果</h2>
    <table>
      <tr>
        <th>阈值</th>
        <th>状态</th>
      </tr>
      ${Object.entries(data.thresholds || {}).map(([name, result]) => `
        <tr>
          <td>${name}</td>
          <td class="${result.ok ? 'success' : 'error'}">${result.ok ? '✓ 通过' : '✗ 失败'}</td>
        </tr>
      `).join('')}
    </table>
  </div>
</body>
</html>
  `;
}