# -*- coding: utf-8 -*-
"""
通用LLM模型性能测试工具
支持多种API平台，只需配置API端点和模型列表
"""

import time
import json
import requests
from typing import List, Optional
import statistics
import logging
from dataclasses import dataclass, asdict, field
from datetime import datetime
from pathlib import Path
from mepybase.easylog import EasyLogManager
import matplotlib.pyplot as plt
import matplotlib

matplotlib.use("Agg")  # 使用非交互式后端

# 设置中文字体
plt.rcParams["font.sans-serif"] = ["Noto Sans CJK SC", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False


@dataclass
class TestConfig:
    """测试配置"""

    api_base: str
    api_key: str
    models: List[str]
    test_prompt: str = "What is the meaning of life? please answer in Chinese."
    timeout: int = 120
    delay_between_requests: int = 1
    output_dir: str = "./test_results"
    log_file: str = "model_tester.log"
    log_level: int = logging.INFO


@dataclass
class TestResult:
    """单个模型测试结果"""

    model: str
    success: bool
    response_time: float
    content: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    error: Optional[str] = None
    http_status: Optional[int] = None
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


class ModelPerfTester:
    """模型性能测试器"""

    def __init__(self, config: TestConfig):
        self.config = config

        # 创建输出目录
        Path(config.output_dir).mkdir(parents=True, exist_ok=True)

        # 初始化日志记录器
        self.logger = EasyLogManager.get_logger(
            name="ModelTester",
            log_file=config.log_file,
            log_folder=config.output_dir,
            level=config.log_level,
        )

        self.results: List[TestResult] = []

        self.logger.info(f"输出目录: {config.output_dir}")
        self.logger.info(f"日志文件: {config.output_dir}/{config.log_file}")

    def test_model(self, model_name: str) -> TestResult:
        """测试单个模型"""
        self.logger.info(f"测试模型: {model_name}...")

        start_time = time.time()
        try:
            headers = {
                "Authorization": f"Bearer {self.config.api_key}",
                "Content-Type": "application/json",
            }

            payload = {
                "model": model_name,
                "messages": [{"role": "user", "content": self.config.test_prompt}],
            }

            response = requests.post(
                f"{self.config.api_base}/chat/completions",
                headers=headers,
                json=payload,
                timeout=self.config.timeout,
            )

            elapsed = time.time() - start_time

            if response.status_code == 200:
                data = response.json()
                content = data["choices"][0]["message"]["content"]
                usage = data.get("usage", {})

                result = TestResult(
                    model=model_name,
                    success=True,
                    response_time=elapsed,
                    content=content[:100] + "..." if len(content) > 100 else content,
                    prompt_tokens=usage.get("prompt_tokens", 0),
                    completion_tokens=usage.get("completion_tokens", 0),
                    total_tokens=usage.get("total_tokens", 0),
                    http_status=200,
                )

                self.logger.info(
                    f"✓ {model_name} 完成 ({elapsed:.2f}s, {result.total_tokens} tokens)"
                )
                # 保存完整的模型输出到日志
                self.logger.debug(f"[{model_name}] 完整输出:\n{content}")

                return result
            else:
                error_msg = response.text[:100]
                result = TestResult(
                    model=model_name,
                    success=False,
                    response_time=elapsed,
                    content="",
                    error=f"HTTP {response.status_code}: {error_msg}",
                    http_status=response.status_code,
                )

                self.logger.warning(
                    f"✗ {model_name} 失败 (HTTP {response.status_code})"
                )
                return result

        except requests.Timeout:
            elapsed = time.time() - start_time
            result = TestResult(
                model=model_name,
                success=False,
                response_time=elapsed,
                content="",
                error=f"Timeout ({self.config.timeout}s)",
            )
            self.logger.error(f"✗ {model_name} 超时")
            return result

        except Exception as e:
            elapsed = time.time() - start_time
            result = TestResult(
                model=model_name,
                success=False,
                response_time=elapsed,
                content="",
                error=str(e),
            )
            self.logger.error(f"✗ {model_name} 错误: {str(e)[:50]}")
            return result

    def run_tests(self) -> List[TestResult]:
        """运行所有测试"""
        self.logger.info(f"开始测试 {len(self.config.models)} 个模型...")
        self.logger.info(f"API地址: {self.config.api_base}")
        self.logger.info(f"测试提示: {self.config.test_prompt}\n")

        for i, model in enumerate(self.config.models, 1):
            self.logger.info(f"[{i}/{len(self.config.models)}]")
            result = self.test_model(model)
            self.results.append(result)

            if i < len(self.config.models):
                time.sleep(self.config.delay_between_requests)

        self.logger.info(f"\n测试完成！总共测试 {len(self.results)} 个模型\n")
        return self.results

    def print_statistics(self):
        """打印统计信息"""
        successful = [r for r in self.results if r.success]
        failed = [r for r in self.results if not r.success]

        print("\n" + "=" * 80)
        print("测试结果摘要")
        print("=" * 80)

        print(f"\n总模型数:   {len(self.results)}")
        print(
            f"成功:       {len(successful)} ({len(successful) / len(self.results) * 100:.1f}%)"
        )
        print(
            f"失败:       {len(failed)} ({len(failed) / len(self.results) * 100:.1f}%)"
        )

        if successful:
            response_times = [r.response_time for r in successful]
            total_tokens = [r.total_tokens for r in successful]

            print("\n" + "-" * 80)
            print("响应时间统计 (秒)")
            print("-" * 80)
            print(f"最快:       {min(response_times):.2f}s")
            print(f"最慢:       {max(response_times):.2f}s")
            print(f"平均:       {statistics.mean(response_times):.2f}s")
            print(f"中位数:     {statistics.median(response_times):.2f}s")
            if len(response_times) > 1:
                print(f"标准差:     {statistics.stdev(response_times):.2f}s")

            print("\n" + "-" * 80)
            print("Token使用统计")
            print("-" * 80)
            print(f"最少:       {min(total_tokens)} tokens")
            print(f"最多:       {max(total_tokens)} tokens")
            print(f"平均:       {statistics.mean(total_tokens):.1f} tokens")
            print(f"中位数:     {statistics.median(total_tokens)} tokens")
            if len(total_tokens) > 1:
                print(f"标准差:     {statistics.stdev(total_tokens):.1f} tokens")

            print("\n" + "-" * 80)
            print("详细结果 (按响应时间排序)")
            print("-" * 80)
            print(f"{'模型':<35} {'时间(s)':>12} {'Tokens':>10} {'效率':>12}")
            print("-" * 80)

            for r in sorted(successful, key=lambda x: x.response_time):
                efficiency = (
                    r.total_tokens / r.response_time if r.response_time > 0 else 0
                )
                print(
                    f"{r.model:<35} {r.response_time:>12.2f} {r.total_tokens:>10} {efficiency:>12.1f}"
                )

            print("\n" + "-" * 80)
            print("Top 5 最快模型")
            print("-" * 80)
            for i, r in enumerate(
                sorted(successful, key=lambda x: x.response_time)[:5], 1
            ):
                efficiency = (
                    r.total_tokens / r.response_time if r.response_time > 0 else 0
                )
                print(
                    f"{i}. {r.model:<30} {r.response_time:.2f}s ({efficiency:.1f} tokens/s)"
                )

            print("\n" + "-" * 80)
            print("Top 5 Token最少模型")
            print("-" * 80)
            for i, r in enumerate(
                sorted(successful, key=lambda x: x.total_tokens)[:5], 1
            ):
                print(f"{i}. {r.model:<30} {r.total_tokens} tokens")

        if failed:
            print("\n" + "-" * 80)
            print("失败的模型")
            print("-" * 80)
            for r in failed:
                status = f"(HTTP {r.http_status})" if r.http_status else ""
                print(f"• {r.model}: {r.error} {status}")

        print("\n" + "=" * 80)

    def save_results(self, filename: str = "test_results.json"):
        """保存测试结果到JSON"""
        filepath = Path(self.config.output_dir) / filename

        results_data = [asdict(r) for r in self.results]

        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(results_data, f, indent=2, ensure_ascii=False)

        self.logger.info(f"结果已保存到: {filepath}")
        return filepath

    def save_report(self, filename: str = "test_report.txt"):
        """保存详细报告到文本文件"""
        filepath = Path(self.config.output_dir) / filename

        successful = [r for r in self.results if r.success]
        failed = [r for r in self.results if not r.success]

        with open(filepath, "w", encoding="utf-8") as f:
            f.write("=" * 80 + "\n")
            f.write("LLM 模型性能测试报告\n")
            f.write("=" * 80 + "\n\n")

            f.write(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"API地址: {self.config.api_base}\n")
            f.write(f"测试模型数: {len(self.results)}\n")
            f.write(f"成功率: {len(successful) / len(self.results) * 100:.1f}%\n\n")

            if successful:
                response_times = [r.response_time for r in successful]
                total_tokens = [r.total_tokens for r in successful]

                f.write("响应时间统计:\n")
                f.write(f"  最快: {min(response_times):.2f}s\n")
                f.write(f"  最慢: {max(response_times):.2f}s\n")
                f.write(f"  平均: {statistics.mean(response_times):.2f}s\n\n")

                f.write("Token统计:\n")
                f.write(f"  最少: {min(total_tokens)} tokens\n")
                f.write(f"  最多: {max(total_tokens)} tokens\n")
                f.write(f"  平均: {statistics.mean(total_tokens):.1f} tokens\n\n")

                f.write("-" * 80 + "\n")
                f.write("详细结果\n")
                f.write("-" * 80 + "\n")

                for r in sorted(successful, key=lambda x: x.response_time):
                    efficiency = (
                        r.total_tokens / r.response_time if r.response_time > 0 else 0
                    )
                    f.write(f"\n模型: {r.model}\n")
                    f.write(f"  响应时间: {r.response_time:.2f}s\n")
                    f.write(f"  Tokens: {r.total_tokens}\n")
                    f.write(f"  效率: {efficiency:.1f} tokens/s\n")

        self.logger.info(f"报告已保存到: {filepath}")
        return filepath

    def plot_metrics(self):
        """绘制性能指标图表"""
        successful = [r for r in self.results if r.success]

        if not successful:
            self.logger.warning("没有成功的测试结果，无法绘制图表")
            return

        # 按响应时间排序（从小到大）
        by_response_time = sorted(successful, key=lambda x: x.response_time)
        # 按token数量排序（从多到少）
        by_tokens = sorted(successful, key=lambda x: x.total_tokens, reverse=True)

        # 创建图表
        fig, axes = plt.subplots(2, 1, figsize=(14, 10))
        fig.suptitle("LLM Model Performance Report", fontsize=16, fontweight="bold")

        # 图1: 按响应时间排序（从小到大）
        ax1 = axes[0]
        # 反转顺序，使得barh中最上面显示响应时间最短的
        by_response_time_reversed = list(reversed(by_response_time))
        models_time = [r.model for r in by_response_time_reversed]
        times = [r.response_time for r in by_response_time_reversed]

        # 计算颜色 - 响应时间越短越绿
        cmap = plt.get_cmap("RdYlGn")
        norm_times = [
            (max(times) - t) / (max(times) - min(times))
            if max(times) > min(times)
            else 0.5
            for t in times
        ]
        colors_time = [cmap(nt) for nt in norm_times]

        bars1 = ax1.barh(models_time, times, color=colors_time)
        ax1.set_xlabel("Response Time (seconds)", fontsize=11, fontweight="bold")
        ax1.set_title(
            "Sorted by Response Time (Ascending)",
            fontsize=12,
            fontweight="bold",
            pad=10,
        )
        ax1.grid(axis="x", alpha=0.3)

        # 在条形图上显示值
        for i, (bar, time) in enumerate(zip(bars1, times)):
            ax1.text(time + 0.5, i, f"{time:.2f}s", va="center", fontsize=9)

        # 图2: 按Token数量排序（从多到少）
        ax2 = axes[1]
        # 反转顺序，使得barh中最上面显示Token最多的
        by_tokens_reversed = list(reversed(by_tokens))
        models_tokens = [r.model for r in by_tokens_reversed]
        token_counts = [r.total_tokens for r in by_tokens_reversed]

        # 计算颜色 - token数越多越绿
        cmap2 = plt.get_cmap("viridis")
        norm_tokens = [
            (t - min(token_counts)) / (max(token_counts) - min(token_counts))
            if max(token_counts) > min(token_counts)
            else 0.5
            for t in token_counts
        ]
        colors_tokens = [cmap2(nt) for nt in norm_tokens]

        bars2 = ax2.barh(models_tokens, token_counts, color=colors_tokens)
        ax2.set_xlabel("Token Count", fontsize=11, fontweight="bold")
        ax2.set_title(
            "Sorted by Token Count (Descending)", fontsize=12, fontweight="bold", pad=10
        )
        ax2.grid(axis="x", alpha=0.3)

        # 在条形图上显示值
        for i, (bar, token_count) in enumerate(zip(bars2, token_counts)):
            ax2.text(token_count + 10, i, f"{token_count}", va="center", fontsize=9)

        # 调整布局
        plt.tight_layout()

        # 保存图表
        filepath = Path(self.config.output_dir) / "performance_metrics.png"
        plt.savefig(filepath, dpi=300, bbox_inches="tight")
        self.logger.info(f"性能指标图表已保存到: {filepath}")
        plt.close()

        return filepath


# ===============================
# 使用示例
# ===============================

if __name__ == "__main__":
    # 配置 iFlow 测试
    config = TestConfig(
        api_base="https://apis.iflow.cn/v1",
        api_key="sk-d2dbd717ab27761fde138535963a09f0",
        models=[
            "iflow-rome-30ba3b",
            "qwen3-coder-plus",
            "qwen3-max",
            "qwen3-vl-plus",
            "kimi-k2-0905",
            "qwen3-max-preview",
            "kimi-k2",
            "deepseek-v3.2",
            "deepseek-r1",
            "deepseek-v3",
            "qwen3-32b",
            "qwen3-235b-a22b-thinking-2507",
            "qwen3-235b-a22b-instruct",
            "qwen3-235b",
        ],
        test_prompt="What is the meaning of life? please answer in Chinese.",
        output_dir="./test_results/iflow",
        log_file="iflow_test.log",
        log_level=logging.INFO,
    )

    # 创建测试器
    tester = ModelPerfTester(config)

    # 运行测试
    results = tester.run_tests()

    # 打印统计信息
    tester.print_statistics()

    # 保存结果
    tester.save_results("test_results.json")
    tester.save_report("test_report.txt")

    # 绘制性能指标图表
    tester.plot_metrics()

    print("\n✓ 所有测试完成！")
    print(f"✓ 结果保存在: {tester.config.output_dir}/")
