#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：performance_executor.py
模块描述：性能测试执行引擎 — 轻量级 asyncio + aiohttp 自研压测引擎

支持三种压测模式：
  - concurrent（固定并发数）
  - qps（固定 QPS）
  - step（阶梯加压）

核心组件：
  - MetricsCollector：线程安全指标采集器，每秒生成快照
  - ResultValidator：断言规则验证器
  - PerformanceExecutor：主引擎，协调 Token 获取、请求发送、指标采集

作者：czj
创建日期：2026-08-12
最后修改：2026-08-12
"""

import asyncio
import json
import logging
import threading
import time
from typing import Any, Callable, Dict, List, Optional

import aiohttp

from app.executors.http_client import HttpClient

logger = logging.getLogger(__name__)


# ==================== MetricsCollector ====================

class MetricsCollector:
    """
    指标采集器

    线程安全，支持每秒采集快照和最终汇总统计。
    使用 threading.Lock 保护共享状态。

    Attributes:
        snapshot_callback: 每秒快照回调函数（用于 WebSocket 推送）
    """

    def __init__(self, snapshot_callback: Optional[Callable[[dict], None]] = None):
        self.snapshot_callback = snapshot_callback
        self._lock = threading.Lock()
        self._start_time = time.time()
        self._reset()

    def _reset(self) -> None:
        """重置周期计数器（每秒重置）"""
        # 周期内（当前秒）
        self._period_success = 0
        self._period_fail = 0
        self._period_times: List[float] = []
        self._period_status_2xx = 0
        self._period_status_3xx = 0
        self._period_status_4xx = 0
        self._period_status_5xx = 0

        # 全局累计
        self.total_requests = 0
        self.success_count = 0
        self.fail_count = 0
        self.all_times: List[float] = []
        self.status_2xx = 0
        self.status_3xx = 0
        self.status_4xx = 0
        self.status_5xx = 0

    def record(self, status: int, elapsed_ms: float, success: bool, error: str = "") -> None:
        """
        记录一次请求结果

        Args:
            status: HTTP 状态码（0 表示网络异常）
            elapsed_ms: 响应时间（毫秒）
            success: 是否成功（断言通过）
            error: 错误描述
        """
        with self._lock:
            self.total_requests += 1
            if success:
                self.success_count += 1
                self._period_success += 1
            else:
                self.fail_count += 1
                self._period_fail += 1

            self._period_times.append(elapsed_ms)
            self.all_times.append(elapsed_ms)

            if 200 <= status < 300:
                self._period_status_2xx += 1
                self.status_2xx += 1
            elif 300 <= status < 400:
                self._period_status_3xx += 1
                self.status_3xx += 1
            elif 400 <= status < 500:
                self._period_status_4xx += 1
                self.status_4xx += 1
            elif status >= 500:
                self._period_status_5xx += 1
                self.status_5xx += 1

    def snapshot(self) -> dict:
        """
        采集当前周期快照并重置周期计数器

        Returns:
            dict: 快照数据（含周期统计和全局累计）
        """
        with self._lock:
            elapsed = time.time() - self._start_time
            period_times = self._period_times
            period_success = self._period_success
            period_fail = self._period_fail

            # 计算周期统计
            tps = period_success / 1.0 if period_success > 0 else 0
            avg = sum(period_times) / len(period_times) if period_times else 0
            p50 = self._percentile(period_times, 50) if period_times else 0
            p90 = self._percentile(period_times, 90) if period_times else 0
            p99 = self._percentile(period_times, 99) if period_times else 0

            snapshot = {
                "elapsed": round(elapsed, 1),
                "tps": round(tps, 2),
                "avg_response_time": round(avg, 2),
                "min_response_time": round(min(period_times), 2) if period_times else 0,
                "max_response_time": round(max(period_times), 2) if period_times else 0,
                "p50": round(p50, 2),
                "p90": round(p90, 2),
                "p99": round(p99, 2),
                "success_count": period_success,
                "fail_count": period_fail,
                "status_2xx": self._period_status_2xx,
                "status_3xx": self._period_status_3xx,
                "status_4xx": self._period_status_4xx,
                "status_5xx": self._period_status_5xx,
                "total_requests": self.total_requests,
                "concurrent_users": 0,  # 由调用者填充
            }

            # 重置周期计数器
            self._period_success = 0
            self._period_fail = 0
            self._period_times = []
            self._period_status_2xx = 0
            self._period_status_3xx = 0
            self._period_status_4xx = 0
            self._period_status_5xx = 0

            return snapshot

    def summary(self) -> dict:
        """
        最终汇总统计

        Returns:
            dict: 完整汇总数据
        """
        times = self.all_times
        elapsed = max(time.time() - self._start_time, 1)
        return {
            "total_requests": self.total_requests,
            "success_count": self.success_count,
            "fail_count": self.fail_count,
            "error_rate": round(self.fail_count / max(self.total_requests, 1) * 100, 2),
            "min_response_time": round(min(times), 2) if times else 0,
            "max_response_time": round(max(times), 2) if times else 0,
            "avg_response_time": round(sum(times) / len(times), 2) if times else 0,
            "p50_response_time": round(self._percentile(times, 50), 2) if times else 0,
            "p90_response_time": round(self._percentile(times, 90), 2) if times else 0,
            "p99_response_time": round(self._percentile(times, 99), 2) if times else 0,
            "actual_tps": round(self.success_count / elapsed, 2),
            "status_2xx": self.status_2xx,
            "status_3xx": self.status_3xx,
            "status_4xx": self.status_4xx,
            "status_5xx": self.status_5xx,
        }

    @staticmethod
    def _percentile(data: List[float], p: float) -> float:
        """
        计算百分位数（线性插值）

        Args:
            data: 排序前的数据列表
            p: 百分位（0-100）

        Returns:
            float: 百分位数值
        """
        if not data:
            return 0.0
        sorted_data = sorted(data)
        k = (len(sorted_data) - 1) * p / 100
        f = int(k)
        c = min(f + 1, len(sorted_data) - 1)
        if f == c:
            return sorted_data[f]
        return sorted_data[f] + (sorted_data[c] - sorted_data[f]) * (k - f)


# ==================== ResultValidator ====================

class ResultValidator:
    """
    结果验证器

    执行断言规则列表，判断一次请求是否符合预期。
    支持三种断言类型：
      - status_code: HTTP 状态码比较
      - response_time: 响应时间比较（ms）
      - response_body: 响应体内容检查（JSON Path / 文本包含）
    """

    @staticmethod
    def validate(assertions: list, status_code: int, body: str, response_time: float) -> bool:
        """
        执行所有断言，全部通过返回 True

        Args:
            assertions: 断言规则列表
            status_code: HTTP 状态码
            body: 响应体文本
            response_time: 响应时间（ms）

        Returns:
            bool: 所有断言通过返回 True
        """
        if not assertions:
            return True

        for assertion in assertions:
            if not ResultValidator._evaluate(assertion, status_code, body, response_time):
                return False
        return True

    @staticmethod
    def _evaluate(assertion: dict, status_code: int, body: str, response_time: float) -> bool:
        """
        评估单条断言

        Args:
            assertion: 断言规则字典
            status_code: HTTP 状态码
            body: 响应体文本
            response_time: 响应时间（ms）

        Returns:
            bool: 断言通过返回 True
        """
        atype = assertion.get("type", "")
        operator = assertion.get("operator", "equals")
        value = assertion.get("value")

        if atype == "status_code":
            return ResultValidator._compare(status_code, operator, value)
        elif atype == "response_time":
            return ResultValidator._compare(response_time, operator, value)
        elif atype == "response_body":
            return ResultValidator._check_body(assertion, body)
        else:
            # 未知断言类型默认通过
            return True

    @staticmethod
    def _compare(actual: Any, operator: str, expected: Any) -> bool:
        """
        比较运算符

        Args:
            actual: 实际值
            operator: 运算符
            expected: 期望值

        Returns:
            bool: 比较结果
        """
        try:
            if operator == "equals":
                return actual == expected
            elif operator == "not_equals":
                return actual != expected
            elif operator == "less_than":
                return float(actual) < float(expected)
            elif operator == "less_or_equal":
                return float(actual) <= float(expected)
            elif operator == "greater_than":
                return float(actual) > float(expected)
            elif operator == "greater_or_equal":
                return float(actual) >= float(expected)
            elif operator == "json_path":
                return ResultValidator._check_json_path(actual, expected)
            return True
        except (TypeError, ValueError) as e:
            logger.warning(f"断言比较异常: {e}")
            return False

    @staticmethod
    def _check_body(assertion: dict, body: str) -> bool:
        """
        检查响应体

        Args:
            assertion: 断言规则（含 path 和 value）
            body: 响应体文本

        Returns:
            bool: 检查通过返回 True
        """
        path = assertion.get("path", "")
        expected = assertion.get("value", "")
        operator = assertion.get("operator", "contains")

        if not path:
            # 无 path 时直接检查文本包含
            if operator == "contains":
                return str(expected) in body
            elif operator == "not_contains":
                return str(expected) not in body
            return True

        # 有 path 时尝试 JSON Path 解析
        try:
            data = json.loads(body)
            actual = data
            for key in path.strip("$.").split("."):
                if isinstance(actual, dict):
                    actual = actual.get(key)
                elif isinstance(actual, list):
                    try:
                        idx = int(key)
                        actual = actual[idx] if 0 <= idx < len(actual) else None
                    except ValueError:
                        actual = None
                else:
                    actual = None
                    break
            return ResultValidator._compare(actual, operator, expected)
        except (json.JSONDecodeError, KeyError, IndexError, TypeError):
            return False

    @staticmethod
    def _check_json_path(actual: Any, expected: Any) -> bool:
        """
        检查 JSON Path 值（占位，预留扩展）

        Args:
            actual: 实际值
            expected: 期望值

        Returns:
            bool: 匹配返回 True
        """
        return str(actual) == str(expected)


# ==================== PerformanceExecutor ====================

class PerformanceExecutor:
    """
    性能测试执行器

    轻量级 asyncio + aiohttp 自研引擎，支持三种压测模式：
      - concurrent: 固定并发数
      - qps: 固定 QPS（每秒请求数）
      - step: 阶梯并发（逐步递增）

    与 PlaywrightExecutor(UI测试) 和 SecurityExecutor(安全测试) 并列。

    Usage:
        >>> executor = PerformanceExecutor(http_client_config)
        >>> result = await executor.execute(task)
    """

    def __init__(self, http_client_config: dict):
        """
        初始化执行器

        Args:
            http_client_config: HttpClient 配置字典
        """
        self.http_client_config = http_client_config
        self._http_client: Optional[HttpClient] = None
        self._token: Optional[str] = None
        self._running = False
        self._metrics: Optional[MetricsCollector] = None
        self._validator = ResultValidator()
        # 登录接口压测时复用的验证码 UUID（首次获取，后续复用）
        self._login_uuid: Optional[str] = None

    def _get_http_client(self) -> HttpClient:
        """获取或创建 HttpClient 实例"""
        if not self._http_client:
            self._http_client = HttpClient(self.http_client_config)
        return self._http_client

    def _ensure_token(self, task) -> str:
        """
        确保有有效的 Token（复用 HttpClient.login）

        Args:
            task: PerformanceTask 任务对象

        Returns:
            str: Bearer Token
        """
        if not self._http_client:
            self._http_client = HttpClient(self.http_client_config)

        if not self._token:
            self._token = self._http_client.login(task.account_key)

        return self._token or ""

    def _is_login_target(self, task) -> bool:
        """
        判断目标任务是否为登录接口

        登录接口需要特殊处理：请求体包含 username/password/code/uuid，
        且请求头需要签名（X-RANDOM/X-TIMESTAMP/X-SIGN/RandomCode）。
        """
        login_path = self.http_client_config.get("auth", {}).get(
            "login_path", "/platform/api/auth/login"
        )
        url = (task.target_url or "").lower()
        return login_path in url

    def _get_login_body(self, task) -> Dict[str, Any]:
        """
        构造登录接口请求体

        使用加密后的密码和验证码 UUID，确保每次压测请求的登录请求合法。

        Args:
            task: PerformanceTask 任务对象

        Returns:
            dict: 登录请求体
        """
        client = self._get_http_client()
        accounts = self.http_client_config.get("accounts", {})
        account = accounts.get(task.account_key, {})
        username = account.get("username", "admin@xty")
        password = account.get("password", "Ubains@13579")
        captcha = accounts.get("captcha", "csba")

        # 加密密码
        encrypted_pwd = HttpClient.encrypt_password(password)

        # 获取验证码 UUID（复用，避免每次压测请求都请求验证码接口）
        if not self._login_uuid:
            self._login_uuid = client.get_captcha_uuid() or ""

        return {
            "username": username,
            "password": encrypted_pwd,
            "code": captcha,
            "uuid": self._login_uuid,
        }

    def _build_headers(self, task) -> Dict[str, str]:
        """
        构造请求头

        Args:
            task: PerformanceTask 任务对象

        Returns:
            dict: 请求头字典
        """
        headers = {}
        if task.headers:
            headers.update(task.headers)

        headers.setdefault("Content-Type", "application/json")
        headers.setdefault("Accept", "application/json, text/plain, */*")

        # 登录接口需要签名（X-RANDOM/X-TIMESTAMP/X-SIGN/RandomCode）
        # 注意：登录接口不能带 Authorization 头，否则签名校验会失败
        # （签名是用 bearer_token="" 生成的，带 Authorization 会导致签名不匹配）
        if self._is_login_target(task):
            body = self._get_login_body(task)
            sign_headers = self._get_http_client()._generate_sign(
                body_data=body, bearer_token=""
            )
            headers.update(sign_headers)
        elif task.auth_required and self._token:
            headers["Authorization"] = f"Bearer {self._token}"

        return headers

    async def execute(self, task, snapshot_callback: Optional[Callable[[dict], None]] = None) -> dict:
        """
        执行性能测试任务

        Args:
            task: PerformanceTask 任务对象
            snapshot_callback: 每秒快照回调（用于 WebSocket 推送）

        Returns:
            dict: 执行结果汇总
        """
        self._metrics = MetricsCollector(snapshot_callback=snapshot_callback)
        self._running = True

        # 1. 获取 Token
        if task.auth_required:
            token = self._ensure_token(task)
            if not token:
                self._running = False
                return {
                    "error": True,
                    "error_message": f"登录失败，无法获取 Token (account: {task.account_key})",
                }

        # 2. 根据 mode 选择执行策略
        try:
            if task.mode == "concurrent":
                result = await self._run_concurrent(task)
            elif task.mode == "qps":
                result = await self._run_qps(task)
            elif task.mode == "step":
                result = await self._run_step(task)
            else:
                result = {"error": True, "error_message": f"未知压测模式: {task.mode}"}
        except Exception as e:
            logger.exception(f"压测执行异常: {e}")
            result = {"error": True, "error_message": str(e)}
        finally:
            self._running = False

        return result

    async def _run_concurrent(self, task) -> dict:
        """
        并发模式：固定并发数，持续压测

        Args:
            task: PerformanceTask 任务对象

        Returns:
            dict: 汇总统计
        """
        semaphore = asyncio.Semaphore(task.concurrency)
        connector = aiohttp.TCPConnector(limit=task.concurrency, limit_per_host=task.concurrency, ssl=False)

        # 启动快照采集协程
        sampler_task = asyncio.create_task(self._snapshot_loop(task))

        try:
            async with aiohttp.ClientSession(connector=connector) as session:
                # 预热阶段：逐步增加并发
                if task.ramp_up > 0:
                    ramp_start = time.time()
                    ramp_workers = []
                    for i in range(task.concurrency):
                        if not self._running:
                            break
                        ramp_workers.append(
                            asyncio.create_task(self._worker(session, task, semaphore))
                        )
                        # 在预热时间内均匀启动
                        await asyncio.sleep(task.ramp_up / max(task.concurrency, 1))
                    ramp_elapsed = time.time() - ramp_start
                    logger.info(f"预热完成: {len(ramp_workers)} 个协程已启动, 耗时 {ramp_elapsed:.1f}s")
                else:
                    # 无预热，直接启动所有协程
                    ramp_workers = [
                        asyncio.create_task(self._worker(session, task, semaphore))
                        for _ in range(task.concurrency)
                    ]

                # 压测阶段：等待指定时长（分段检查停止）
                await self._sleep_with_check(task.duration)
                self._running = False

                # 等待所有 worker 结束
                await asyncio.gather(*ramp_workers, return_exceptions=True)
        finally:
            sampler_task.cancel()
            try:
                await sampler_task
            except asyncio.CancelledError:
                pass

        return self._metrics.summary()

    async def _run_qps(self, task) -> dict:
        """
        QPS 模式：控制每秒请求速率

        通过令牌桶算法控制请求速率，确保每秒请求数不超过 qps_target。

        Args:
            task: PerformanceTask 任务对象

        Returns:
            dict: 汇总统计
        """
        connector = aiohttp.TCPConnector(limit=200, limit_per_host=200, ssl=False)
        sampler_task = asyncio.create_task(self._snapshot_loop(task))

        # 令牌桶参数
        rate = task.qps_target  # 每秒令牌数
        tokens = rate           # 初始令牌数
        last_refill = time.time()
        min_interval = 1.0 / max(rate, 1)  # 最小请求间隔

        # 并发请求控制
        semaphore = asyncio.Semaphore(task.concurrency if task.concurrency > 0 else 50)

        try:
            async with aiohttp.ClientSession(connector=connector) as session:
                end_time = time.time() + task.duration
                pending = []

                while self._running and time.time() < end_time:
                    # 补充令牌
                    now = time.time()
                    elapsed = now - last_refill
                    tokens = min(rate, tokens + elapsed * rate)
                    last_refill = now

                    # 消耗令牌并发起请求
                    while tokens >= 1.0 and self._running:
                        tokens -= 1.0
                        task_obj = asyncio.create_task(
                            self._send_request(session, task)
                        )
                        pending.append(task_obj)
                        await asyncio.sleep(min_interval * 0.5)  # 微间隔避免突发

                    # 清理已完成的任务
                    pending = [p for p in pending if not p.done()]

                    # 等待一小段时间再检查令牌（分段检查停止）
                    if tokens < 1.0:
                        await self._sleep_with_check(min_interval * 0.5)

                self._running = False

                # 等待剩余请求完成
                if pending:
                    await asyncio.gather(*pending, return_exceptions=True)
        finally:
            sampler_task.cancel()
            try:
                await sampler_task
            except asyncio.CancelledError:
                pass

        return self._metrics.summary()

    async def _run_step(self, task) -> dict:
        """
        阶梯模式：逐步递增并发数

        step_concurrency = [10, 50, 100, 200] 表示 4 个阶梯
        step_duration = 60 表示每阶梯持续 60 秒

        Args:
            task: PerformanceTask 任务对象

        Returns:
            dict: 汇总统计
        """
        step_concurrency = task.step_concurrency or [task.concurrency]
        step_duration = task.step_duration or task.duration

        if not step_concurrency:
            step_concurrency = [task.concurrency]

        connector = aiohttp.TCPConnector(limit=200, limit_per_host=200, ssl=False)
        sampler_task = asyncio.create_task(self._snapshot_loop(task))

        try:
            async with aiohttp.ClientSession(connector=connector) as session:
                for step_idx, concurrent in enumerate(step_concurrency):
                    if not self._running:
                        break

                    logger.info(f"阶梯 {step_idx + 1}/{len(step_concurrency)}: 并发数={concurrent}, 时长={step_duration}s")
                    semaphore = asyncio.Semaphore(concurrent)

                    # 启动当前阶梯的 workers
                    workers = [
                        asyncio.create_task(self._worker(session, task, semaphore))
                        for _ in range(concurrent)
                    ]

                    # 持续当前阶梯时长（分段检查停止）
                    await self._sleep_with_check(step_duration)

                    # 停止当前阶梯的 workers
                    self._running = False
                    await asyncio.gather(*workers, return_exceptions=True)
                    self._running = True  # 为下一阶梯重置

        finally:
            self._running = False
            sampler_task.cancel()
            try:
                await sampler_task
            except asyncio.CancelledError:
                pass

        return self._metrics.summary()

    async def _worker(self, session: aiohttp.ClientSession, task, semaphore: asyncio.Semaphore) -> None:
        """
        工作协程：循环发送请求直到停止

        Args:
            session: aiohttp 会话
            task: PerformanceTask 任务对象
            semaphore: 并发信号量
        """
        while self._running:
            async with semaphore:
                await self._send_request(session, task)

    async def _send_request(self, session: aiohttp.ClientSession, task) -> None:
        """
        发送单个 HTTP 请求并记录指标

        Args:
            session: aiohttp 会话
            task: PerformanceTask 任务对象
        """
        start = time.time()
        status = 0
        body = ""
        elapsed_ms = 0.0

        try:
            headers = self._build_headers(task)
            method = task.method.upper()
            url = task.target_url

            # 登录接口使用自动构造的请求体（含加密密码 + 验证码UUID）
            if self._is_login_target(task):
                json_data = self._get_login_body(task)
            else:
                json_data = task.body if method in ("POST", "PUT", "PATCH") else None

            async with session.request(
                method, url,
                headers=headers,
                json=json_data,
                ssl=False,
                timeout=aiohttp.ClientTimeout(total=30),
            ) as resp:
                status = resp.status
                body = await resp.text()
                elapsed_ms = (time.time() - start) * 1000

                # 执行断言
                assert_ok = self._validator.validate(
                    task.assertions, status, body, elapsed_ms
                )

                # 记录指标
                self._metrics.record(status, elapsed_ms, assert_ok)

        except asyncio.TimeoutError:
            elapsed_ms = (time.time() - start) * 1000
            self._metrics.record(0, elapsed_ms, False, "timeout")
        except aiohttp.ClientError as e:
            elapsed_ms = (time.time() - start) * 1000
            self._metrics.record(0, elapsed_ms, False, str(e))
        except Exception as e:
            elapsed_ms = (time.time() - start) * 1000
            self._metrics.record(0, elapsed_ms, False, str(e))

    async def _sleep_with_check(self, seconds: float) -> None:
        """
        分段睡眠，期间持续检查停止标志，保证 stop() 能快速生效

        Args:
            seconds: 总睡眠时长（秒）
        """
        remaining = float(seconds)
        while remaining > 0 and self._running:
            # 分 0.5s 一段检查，停止信号最长 0.5s 内响应
            chunk = min(remaining, 0.5)
            await asyncio.sleep(chunk)
            remaining -= chunk

    async def _snapshot_loop(self, task) -> None:
        """
        快照采集循环：每秒采集一次指标并回调

        Args:
            task: PerformanceTask 任务对象
        """
        while self._running:
            await asyncio.sleep(1)
            try:
                snapshot = self._metrics.snapshot()
                if self._metrics.snapshot_callback:
                    self._metrics.snapshot_callback(snapshot)
            except Exception as e:
                logger.error(f"快照采集异常: {e}")

    def stop(self) -> None:
        """停止压测"""
        self._running = False
        logger.info("性能测试执行器已停止")