// Sub2API Gateway Performance Test
// Gateway API 性能测试

import http from 'k6/http';
import { check, sleep, group } from 'k6';
import { Rate, Trend, Counter, Gauge } from 'k6/metrics';
import { config, getBaseUrl, getAuthToken } from '../common/utils.js';
import { httpGet, httpPost, randomSleep, generateTestAPIKeyName } from '../common/utils.js';

// ============= 自定义指标 =============

// Gateway 请求指标
const gatewayRequestDuration = new Trend('gateway_request_duration');
const gatewayTTFT = new Trend('gateway_ttft');  // Time To First Token
const gatewayTokenThroughput = new Trend('gateway_token_throughput');
const gatewayErrorRate = new Rate('gateway_errors');

// 平台分类指标
const platformMetrics = {
  openai: {
    duration: new Trend('gateway_openai_duration'),
    errors: new Rate('gateway_openai_errors'),
  },
  claude: {
    duration: new Trend('gateway_claude_duration'),
    errors: new Rate('gateway_claude_errors'),
  },
  gemini: {
    duration: new Trend('gateway_gemini_duration'),
    errors: new Rate('gateway_gemini_errors'),
  },
};

// ============= 测试配置 =============

export const options = {
  scenarios: {
    gateway_load: {
      executor: 'ramping-vus',
      startVUs: 5,
      stages: [
        { duration: '2m', target: 10 },   // 预热
        { duration: '5m', target: 50 },   // 正常负载
        { duration: '2m', target: 100 },  // 峰值负载
        { duration: '2m', target: 0 },    // 冷却
      ],
    },
  },

  thresholds: {
    // Gateway 整体阈值
    'gateway_request_duration': ['p(95)<2000', 'p(99)<5000'],
    'gateway_errors': ['rate<0.05'],  // 5% 以下错误率

    // TTFT 阈值
    'gateway_ttft': ['p(95)<3000', 'p(99)<5000'],

    // 按平台分类
    'gateway_openai_duration': ['p(95)<1500'],
    'gateway_claude_duration': ['p(95)<2000'],
    'gateway_gemini_duration': ['p(95)<1500'],
  },
};

// ============= 测试数据准备 =============

// 获取测试用 API Key
function getTestAPIKey() {
  const token = getAuthToken();

  // 列出 API Keys 获取第一个
  const res = http.get(`${getBaseUrl()}/api/v1/keys`, {
    headers: {
      'Authorization': `Bearer ${token}`,
      'Content-Type': 'application/json',
    },
    tags: { name: 'list_api_keys' },
  });

  if (res.status === 200) {
    const keys = res.json('data');
    if (keys && keys.length > 0) {
      return keys[0].key;
    }
  }

  // 如果没有 API Key，创建一个
  const createRes = http.post(
    `${getBaseUrl()}/api/v1/keys`,
    JSON.stringify({
      name: `perf-test-${Date.now()}`,
    }),
    {
      headers: {
        'Authorization': `Bearer ${token}`,
        'Content-Type': 'application/json',
      },
      tags: { name: 'create_api_key' },
    }
  );

  if (createRes.status === 201) {
    return createRes.json('key');
  }

  throw new Error('Failed to get or create API key for testing');
}

// ============= 测试场景 =============

/**
 * 测试 OpenAI Chat Completions
 */
function testOpenAIChat(apiKey) {
  const start = Date.now();

  const res = http.post(
    `${getBaseUrl()}/v1/chat/completions`,
    JSON.stringify({
      model: 'gpt-3.5-turbo',
      messages: [
        { role: 'user', content: 'Say hello in one sentence.' }
      ],
      max_tokens: 50,
      temperature: 0.7,
    }),
    {
      headers: {
        'Authorization': `Bearer ${apiKey}`,
        'Content-Type': 'application/json',
      },
      tags: { name: 'gateway_openai_chat' },
    }
  );

  const duration = Date.now() - start;

  // 记录指标
  platformMetrics.openai.duration.add(duration);
  platformMetrics.openai.errors.add(res.status !== 200);

  check(res, {
    'OpenAI Chat: status is 200': (r) => r.status === 200,
    'OpenAI Chat: has content': (r) => r.json('choices[0].message.content') !== undefined,
  });

  return res;
}

/**
 * 测试 OpenAI Streaming
 */
function testOpenAIStream(apiKey) {
  const start = Date.now();
  let firstTokenTime = 0;
  let tokenCount = 0;

  const res = http.post(
    `${getBaseUrl()}/v1/chat/completions`,
    JSON.stringify({
      model: 'gpt-3.5-turbo',
      messages: [
        { role: 'user', content: 'Count from 1 to 5.' }
      ],
      max_tokens: 100,
      stream: true,
    }),
    {
      headers: {
        'Authorization': `Bearer ${apiKey}`,
        'Content-Type': 'application/json',
      },
      tags: { name: 'gateway_openai_stream' },
    }
  );

  const duration = Date.now() - start;

  // 解析 SSE 流获取 TTFT
  if (res.headers['Content-Type']?.includes('text/event-stream')) {
    const lines = res.body.split('\n');
    for (const line of lines) {
      if (line.startsWith('data: ') && !line.includes('[DONE]')) {
        if (firstTokenTime === 0) {
          firstTokenTime = Date.now();
        }
        tokenCount++;
      }
    }
  }

  const ttft = firstTokenTime > 0 ? firstTokenTime - start : duration;

  // 记录指标
  platformMetrics.openai.duration.add(duration);
  platformMetrics.openai.errors.add(res.status !== 200);
  gatewayTTFT.add(ttft);

  if (tokenCount > 0) {
    gatewayTokenThroughput.add(tokenCount / (duration / 1000));
  }

  check(res, {
    'OpenAI Stream: status is 200': (r) => r.status === 200,
    'OpenAI Stream: is streaming': (r) => r.headers['Content-Type']?.includes('text/event-stream'),
  });

  return res;
}

/**
 * 测试 Claude Messages API
 */
function testClaudeMessages(apiKey) {
  const start = Date.now();

  const res = http.post(
    `${getBaseUrl()}/v1/messages`,
    JSON.stringify({
      model: 'claude-3-5-haiku-20241022',
      max_tokens: 50,
      messages: [
        { role: 'user', content: 'Say hello in one sentence.' }
      ],
    }),
    {
      headers: {
        'Authorization': `Bearer ${apiKey}`,
        'Content-Type': 'application/json',
        'x-api-key': apiKey,
        'anthropic-version': '2023-06-01',
      },
      tags: { name: 'gateway_claude_messages' },
    }
  );

  const duration = Date.now() - start;

  // 记录指标
  platformMetrics.claude.duration.add(duration);
  platformMetrics.claude.errors.add(res.status !== 200);

  check(res, {
    'Claude Messages: status is 200': (r) => r.status === 200,
    'Claude Messages: has content': (r) => r.json('content[0].text') !== undefined,
  });

  return res;
}

/**
 * 测试 Gemini Generate Content
 */
function testGeminiGenerate(apiKey) {
  const start = Date.now();

  const res = http.post(
    `${getBaseUrl()}/v1beta/models/gemini-pro:generateContent`,
    JSON.stringify({
      contents: [
        {
          parts: [{ text: 'Say hello in one sentence.' }]
        }
      ],
      generationConfig: {
        maxOutputTokens: 50,
      },
    }),
    {
      headers: {
        'Authorization': `Bearer ${apiKey}`,
        'Content-Type': 'application/json',
      },
      tags: { name: 'gateway_gemini_generate' },
    }
  );

  const duration = Date.now() - start;

  // 记录指标
  platformMetrics.gemini.duration.add(duration);
  platformMetrics.gemini.errors.add(res.status !== 200);

  check(res, {
    'Gemini Generate: status is 200': (r) => r.status === 200,
    'Gemini Generate: has content': (r) => r.json('candidates[0].content.parts[0].text') !== undefined,
  });

  return res;
}

/**
 * 综合 Gateway 测试
 */
function testGatewayMixed(apiKey) {
  // 模拟真实用户行为：70% 非流式 + 30% 流式
  const rand = Math.random();

  if (rand < 0.3) {
    return testOpenAIStream(apiKey);
  } else if (rand < 0.5) {
    return testOpenAIChat(apiKey);
  } else if (rand < 0.7) {
    return testClaudeMessages(apiKey);
  } else {
    return testGeminiGenerate(apiKey);
  }
}

// ============= 主测试函数 =============

export default function () {
  // 提前获取 API Key（每个 VU 一次）
  const apiKey = __ITER__ === 0 ? getTestAPIKey() : null;

  // 或者使用全局缓存
  if (!globalThis.__testApiKey__) {
    globalThis.__testApiKey__ = getTestAPIKey();
  }

  group('Gateway API Performance', () => {
    // 随机选择测试场景
    const testType = __VU % 4;

    switch (testType) {
      case 0:
        testOpenAIChat(globalThis.__testApiKey__);
        break;
      case 1:
        testOpenAIStream(globalThis.__testApiKey__);
        break;
      case 2:
        testClaudeMessages(globalThis.__testApiKey__);
        break;
      case 3:
        testGeminiGenerate(globalThis.__testApiKey__);
        break;
      default:
        testGatewayMixed(globalThis.__testApiKey__);
    }
  });

  // 模拟用户思考时间
  randomSleep(0.5, 2);
}

// ============= 测试结束清理 =============

export function handleSummary(data) {
  return {
    'gateway-performance-report.json': JSON.stringify(data, null, 2),
    'gateway-performance-summary.txt': generateSummary(data),
  };
}

function generateSummary(data) {
  const metrics = data.metrics;

  return `
================================================================================
                    Sub2API Gateway 性能测试报告
================================================================================

测试时间: ${new Date().toISOString()}
测试持续: ${(data.state.testRunDurationMs / 1000 / 60).toFixed(2)} 分钟
峰值 VU: ${data.state.metrics?.vus?.peak || 0}

--------------------------------------------------------------------------------
                            核心指标
--------------------------------------------------------------------------------

总请求数:     ${metrics?.requests_total?.values?.count || 0}
成功请求:     ${metrics?.requests_total?.values?.passes || 0}
失败请求:     ${metrics?.requests_total?.values?.fails || 0}
错误率:       ${((metrics?.gateway_errors?.values?.rate || 0) * 100).toFixed(2)}%

平均响应时间: ${metrics?.gateway_request_duration?.values?.avg?.toFixed(2) || 0} ms
P50 响应时间: ${metrics?.gateway_request_duration?.values?.['p(50)']?.toFixed(2) || 0} ms
P95 响应时间: ${metrics?.gateway_request_duration?.values?.['p(95)']?.toFixed(2) || 0} ms
P99 响应时间: ${metrics?.gateway_request_duration?.values?.['p(99)']?.toFixed(2) || 0} ms
最大响应时间: ${metrics?.gateway_request_duration?.values?.max?.toFixed(2) || 0} ms

--------------------------------------------------------------------------------
                          TTFT (Time To First Token)
--------------------------------------------------------------------------------

平均 TTFT:    ${metrics?.gateway_ttft?.values?.avg?.toFixed(2) || 0} ms
P95 TTFT:     ${metrics?.gateway_ttft?.values?.['p(95)']?.toFixed(2) || 0} ms
P99 TTFT:     ${metrics?.gateway_ttft?.values?.['p(99)']?.toFixed(2) || 0} ms

--------------------------------------------------------------------------------
                          按平台分类
--------------------------------------------------------------------------------

OpenAI:
  平均响应时间: ${metrics?.gateway_openai_duration?.values?.avg?.toFixed(2) || 0} ms
  P95 响应时间: ${metrics?.gateway_openai_duration?.values?.['p(95)']?.toFixed(2) || 0} ms
  错误率:       ${((metrics?.gateway_openai_errors?.values?.rate || 0) * 100).toFixed(2)}%

Claude:
  平均响应时间: ${metrics?.gateway_claude_duration?.values?.avg?.toFixed(2) || 0} ms
  P95 响应时间: ${metrics?.gateway_claude_duration?.values?.['p(95)']?.toFixed(2) || 0} ms
  错误率:       ${((metrics?.gateway_claude_errors?.values?.rate || 0) * 100).toFixed(2)}%

Gemini:
  平均响应时间: ${metrics?.gateway_gemini_duration?.values?.avg?.toFixed(2) || 0} ms
  P95 响应时间: ${metrics?.gateway_gemini_duration?.values?.['p(95)']?.toFixed(2) || 0} ms
  错误率:       ${((metrics?.gateway_gemini_errors?.values?.rate || 0) * 100).toFixed(2)}%

--------------------------------------------------------------------------------
                            阈值检查
--------------------------------------------------------------------------------

${checkThresholds(data)}

================================================================================
                              测试完成
================================================================================
`;
}

function checkThresholds(data) {
  const checks = [];
  const thresholds = options.thresholds;

  for (const [metric, threshold] of Object.entries(thresholds)) {
    const actual = data.metrics?.[metric]?.values;
    if (!actual) continue;

    const p95 = actual['p(95)'];
    const thresholdValue = parseFloat(threshold[0]);

    if (p95 !== undefined) {
      const passed = p95 <= thresholdValue;
      checks.push(`  ${passed ? 'OK' : 'FAIL'} ${metric}: P95=${p95.toFixed(2)}ms (threshold: ${thresholdValue}ms)`);
    }
  }

  return checks.join('\n');
}