#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：performance_service.py
模块描述：性能测试服务层，负责 CRUD、执行调度、快照采集与 WebSocket 推送

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

import asyncio
import json
import logging
import time
from datetime import datetime
from types import SimpleNamespace
from typing import Any, Callable, Dict, List, Optional, Tuple

from sqlalchemy import select, func, desc
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.performance import PerformanceTask, PerformanceSnapshot
from app.models.performance_output import PerformanceTaskOutput
from app.models.api_preset import ApiPreset
from app.schemas.performance import PerformanceTaskCreate
from app.executors.performance_executor import PerformanceExecutor
from app.websocket.manager import manager
from app.utils.id_generator import generate_id

logger = logging.getLogger(__name__)

# 默认的系统登录配置（与 playwright_executor.LOGIN_CONFIG / login_template_service 一致）
_DEFAULT_BASE_URL = "https://192.168.5.44"
_DEFAULT_ACCOUNTS = {
    "superadmin": {"username": "admin@xty", "password": "Ubains@13579"},
    "admin": {"username": "admin@xty", "password": "Ubains@13579"},
    "user": {"username": "admin@xty", "password": "Ubains@13579"},
    "captcha": "csba",
}
_DEFAULT_AUTH_CONFIG = {
    "login_path": "/platform/api/auth/login",
    "captcha_path": "/platform/api/code",
}


class PerformanceService:
    """
    性能测试服务

    提供性能测试任务的 CRUD 操作、执行调度、实时监控和报告生成。
    执行引擎在后台线程池中运行，避免阻塞 FastAPI 事件循环。
    """

    # 正在运行的任务缓存：{task_id: PerformanceExecutor}
    _running_executors: Dict[str, PerformanceExecutor] = {}

    def __init__(self, db: AsyncSession):
        """
        初始化服务

        Args:
            db: 数据库会话
        """
        self.db = db

    # ==================== CRUD ====================

    async def get_task(self, task_id: str) -> Optional[PerformanceTask]:
        """
        获取任务详情

        Args:
            task_id: 任务ID

        Returns:
            PerformanceTask: 任务对象，不存在返回 None
        """
        return await self.db.get(PerformanceTask, task_id)

    async def list_tasks(
        self,
        page: int = 1,
        page_size: int = 20,
        status: Optional[str] = None,
    ) -> Tuple[List[PerformanceTask], int]:
        """
        获取任务列表（分页）

        Args:
            page: 页码
            page_size: 每页数量
            status: 按状态过滤

        Returns:
            tuple: (任务列表, 总数)
        """
        query = select(PerformanceTask)
        count_query = select(func.count()).select_from(PerformanceTask)

        if status:
            query = query.where(PerformanceTask.status == status)
            count_query = count_query.where(PerformanceTask.status == status)

        total = (await self.db.execute(count_query)).scalar() or 0

        query = (
            query.order_by(desc(PerformanceTask.created_at))
            .offset((page - 1) * page_size)
            .limit(page_size)
        )
        result = await self.db.execute(query)
        items = list(result.scalars().all())

        return items, total

    async def create_task(self, data: PerformanceTaskCreate) -> PerformanceTask:
        """
        创建性能测试任务

        如果传入了 preset_id，从接口预设继承接口配置字段
        （method/target_url/headers/body/body_template/capture_rules/
        auth_required/sign_request/account_key），
        任务显式传入的字段覆盖预设值。

        Args:
            data: 任务创建 Schema（含 preset_id 时只继承未显式传入的接口字段）

        Returns:
            PerformanceTask: 创建的任务对象
        """
        # 用户显式传入的字段（排除 schema 默认值，避免默认值覆盖预设）
        explicit = data.model_dump(exclude_unset=True)

        # 从接口预设继承字段
        preset_id = explicit.get("preset_id")
        if preset_id:
            preset = await self.db.get(ApiPreset, preset_id)
            if not preset:
                raise ValueError(f"接口预设不存在: {preset_id}")
            explicit = self._inherit_from_preset(explicit, preset)
            if not explicit.get("target_url"):
                raise ValueError("目标接口 URL 不能为空")

        # 用 schema 默认值补齐其余压测参数
        payload = PerformanceTaskCreate(**explicit).model_dump()

        assertions = payload.pop("assertions", [])
        task = PerformanceTask(
            id=generate_id("perf"),
            **payload,
        )
        task.assertions = assertions
        self.db.add(task)
        await self.db.commit()
        await self.db.refresh(task)
        logger.info(f"创建性能测试任务: {task.id} - {task.name}")
        return task

    @staticmethod
    def _inherit_from_preset(data: Dict[str, Any], preset: ApiPreset) -> Dict[str, Any]:
        """
        从接口预设继承字段（任务显式传入的值覆盖预设默认值）

        输入的 data 为 exclude_unset 后的显式字段字典，
        仅填充未显式传入的接口配置字段。

        Args:
            data: 任务显式传入的字段（exclude_unset 结果）
            preset: 接口预设对象

        Returns:
            dict: 合并后的任务数据
        """
        # 接口配置字段：任务显式传入的优先，否则用预设值
        if not data.get("method"):
            data["method"] = preset.method
        if not data.get("target_url"):
            data["target_url"] = preset.target_url
        if not data.get("headers"):
            data["headers"] = preset.headers or {}
        if data.get("body") is None:
            data["body"] = preset.body
        if not data.get("body_template"):
            data["body_template"] = preset.body_template
        if data.get("capture_rules") is None:
            data["capture_rules"] = preset.capture_rules or []
        if data.get("auth_required") is None:
            data["auth_required"] = preset.auth_required
        if data.get("sign_request") is None:
            data["sign_request"] = preset.sign_request
        if not data.get("account_key"):
            data["account_key"] = preset.account_key
        return data

    @staticmethod
    def _inherit_from_preset(data: Dict[str, Any], preset: ApiPreset) -> Dict[str, Any]:
        """
        从接口预设继承字段（任务显式传入的值覆盖预设默认值）

        Args:
            data: 任务创建数据（model_dump 后含全部字段）
            preset: 接口预设对象

        Returns:
            dict: 合并后的任务数据
        """
        # 预设字段作为默认值，任务显式传入的字段覆盖
        data["method"] = data.get("method") or preset.method
        data["target_url"] = data.get("target_url") or preset.target_url
        data["headers"] = data.get("headers") or (preset.headers or {})
        data["body"] = data["body"] if data.get("body") is not None else preset.body
        data["body_template"] = data.get("body_template") or preset.body_template
        data["capture_rules"] = (
            data.get("capture_rules")
            if data.get("capture_rules") is not None
            else (preset.capture_rules or [])
        )
        data["auth_required"] = (
            data.get("auth_required")
            if data.get("auth_required") is not None
            else preset.auth_required
        )
        data["sign_request"] = (
            data.get("sign_request")
            if data.get("sign_request") is not None
            else preset.sign_request
        )
        data["account_key"] = data.get("account_key") or preset.account_key
        return data

    async def update_task(
        self, task_id: str, data: Dict[str, Any]
    ) -> Optional[PerformanceTask]:
        """
        更新性能测试任务

        Args:
            task_id: 任务ID
            data: 更新数据

        Returns:
            PerformanceTask: 更新后的任务，不存在返回 None
        """
        task = await self.get_task(task_id)
        if not task:
            return None

        # 不允许更新正在运行的任务
        if task.status == "running":
            raise ValueError("正在运行的任务不能修改")

        assertions = data.pop("assertions", None)
        for key, value in data.items():
            if value is not None and hasattr(task, key):
                setattr(task, key, value)
        if assertions is not None:
            task.assertions = assertions

        await self.db.commit()
        await self.db.refresh(task)
        return task

    async def delete_task(self, task_id: str) -> bool:
        """
        删除性能测试任务

        Args:
            task_id: 任务ID

        Returns:
            bool: 是否成功删除
        """
        task = await self.get_task(task_id)
        if not task:
            return False

        if task.status == "running":
            raise ValueError("正在运行的任务不能删除")

        await self.db.delete(task)
        await self.db.commit()
        return True

    # ==================== 执行调度 ====================

    async def run_task(
        self,
        task_id: str,
        config_override: Optional[Dict[str, Any]] = None,
        progress_callback: Optional[Callable] = None,
    ) -> Dict[str, Any]:
        """
        执行性能测试任务

        在后台线程池中运行，避免阻塞 FastAPI 事件循环。
        执行完成后自动更新任务状态和统计结果。

        Args:
            task_id: 任务ID
            config_override: 可覆盖的配置项（如临时修改并发数）
            progress_callback: 进度回调

        Returns:
            dict: 任务数据（含执行结果）
        """
        task = await self.get_task(task_id)
        if not task:
            raise ValueError(f"任务不存在: {task_id}")

        if task.status == "running":
            raise ValueError("任务正在执行中，不能重复启动")

        # 应用配置覆盖
        if config_override:
            for key, value in config_override.items():
                if hasattr(task, key):
                    setattr(task, key, value)

        # 更新状态为运行中
        task.status = "running"
        task.start_time = datetime.utcnow()
        task.error_message = None
        await self.db.flush()

        # 预先把 ORM 对象转成纯字典，避免跨线程 SQLAlchemy 脱钩问题
        task_dict = task.to_dict()

        # 快照收集缓冲区（跨线程收集，执行完成后批量写入）
        snapshot_buffer: List[Dict[str, Any]] = []
        ws_group = f"perf_{task_id}"

        def snapshot_callback(snapshot: dict) -> None:
            """每秒快照回调：收集快照 + WebSocket 推送"""
            snapshot["task_id"] = task_id
            snapshot["timestamp"] = time.time()
            snapshot_buffer.append(dict(snapshot))

            # WebSocket 推送（在回调线程中执行）
            try:
                ws_msg = {
                    "type": "perf_snapshot",
                    "task_id": task_id,
                    "data": snapshot,
                }
                # 使用 asyncio.run_coroutine_threadsafe 在事件循环中广播
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    asyncio.run_coroutine_threadsafe(
                        manager.broadcast(ws_group, ws_msg), loop
                    )
            except Exception as e:
                logger.warning(f"WebSocket 推送快照失败: {e}")

        try:
            # 在线程池中执行
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None,
                self._run_sync,
                task_dict,
                snapshot_callback,
            )

            if result.get("error"):
                task.status = "failed"
                task.error_message = result.get("error_message", "执行异常")
                logger.error(f"性能测试执行失败 [{task_id}]: {result.get('error_message')}")
            else:
                task.status = "completed"
                # 更新统计字段
                self._apply_result(task, result)
                logger.info(f"性能测试执行完成 [{task_id}]: "
                            f"请求={result.get('total_requests', 0)}, "
                            f"TPS={result.get('actual_tps', 0)}")

        except Exception as e:
            task.status = "failed"
            task.error_message = str(e)
            logger.exception(f"性能测试执行异常 [{task_id}]: {e}")

        finally:
            task.end_time = datetime.utcnow()
            if task.start_time:
                task.duration_actual = (task.end_time - task.start_time).total_seconds()
            else:
                task.duration_actual = 0.0
            await self.db.flush()

            # 批量保存快照
            await self._save_snapshots(task_id, snapshot_buffer)

            # 保存响应捕获输出（供其他任务 REF 引用）
            await self._save_outputs(task_id, result.get("captured") if result else None)

            await self.db.commit()

            # 清理运行缓存
            self._running_executors.pop(task_id, None)

            # 发送完成通知
            try:
                asyncio.ensure_future(
                    manager.broadcast(ws_group, {
                        "type": "perf_complete",
                        "task_id": task_id,
                        "status": task.status,
                        "data": task.to_dict(),
                    })
                )
            except Exception:
                pass

        return task.to_dict()

    def _run_sync(
        self,
        task_dict: Dict[str, Any],
        snapshot_callback: Callable,
    ) -> Dict[str, Any]:
        """
        同步执行包装器（在后台线程运行）

        创建独立事件循环运行 asyncio 引擎，避免与 FastAPI 事件循环冲突。

        Args:
            task_dict: 任务数据字典
            snapshot_callback: 快照回调

        Returns:
            dict: 执行结果汇总
        """
        # 设置独立事件循环（参考 PlaywrightExecutor / SecurityExecutor 模式）
        asyncio.set_event_loop(None)
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            # 构建 http_client_config
            client_config = self._build_client_config(task_dict)

            # 创建 SimpleNamespace 以支持属性访问（executor 使用 task.mode 等）
            task_obj = SimpleNamespace(**task_dict)

            executor = PerformanceExecutor(client_config, ref_provider=self._build_ref_provider())
            self._running_executors[task_dict.get("id", "")] = executor

            result = loop.run_until_complete(
                executor.execute(task_obj, snapshot_callback=snapshot_callback)
            )
            return result
        except Exception as e:
            logger.exception(f"_run_sync 异常: {e}")
            return {"error": True, "error_message": str(e)}
        finally:
            loop.close()
            # 清理运行缓存
            tid = task_dict.get("id", "")
            self._running_executors.pop(tid, None)

    async def stop_task(self, task_id: str) -> bool:
        """
        停止正在运行的性能测试任务

        Args:
            task_id: 任务ID

        Returns:
            bool: 是否成功停止
        """
        executor = self._running_executors.get(task_id)
        if not executor:
            logger.warning(f"未找到正在运行的任务: {task_id}")
            return False

        executor.stop()
        logger.info(f"已发送停止信号: {task_id}")

        # 更新数据库状态
        task = await self.get_task(task_id)
        if task and task.status == "running":
            task.status = "cancelled"
            task.end_time = datetime.utcnow()
            if task.start_time:
                task.duration_actual = (task.end_time - task.start_time).total_seconds()
            await self.db.commit()

        return True

    # ==================== 快照与报告 ====================

    async def get_snapshots(
        self,
        task_id: str,
        page: int = 1,
        page_size: int = 500,
    ) -> Tuple[List[PerformanceSnapshot], int]:
        """
        获取任务快照数据

        Args:
            task_id: 任务ID
            page: 页码
            page_size: 每页数量

        Returns:
            tuple: (快照列表, 总数)
        """
        query = (
            select(PerformanceSnapshot)
            .where(PerformanceSnapshot.task_id == task_id)
            .order_by(PerformanceSnapshot.elapsed.asc())
        )
        count_query = (
            select(func.count())
            .select_from(PerformanceSnapshot)
            .where(PerformanceSnapshot.task_id == task_id)
        )

        total = (await self.db.execute(count_query)).scalar() or 0
        query = query.offset((page - 1) * page_size).limit(page_size)
        result = await self.db.execute(query)
        snapshots = list(result.scalars().all())

        return snapshots, total

    async def get_report(self, task_id: str) -> Optional[Dict[str, Any]]:
        """
        获取性能测试报告

        Args:
            task_id: 任务ID

        Returns:
            dict: 报告数据（含摘要、指标、快照），任务不存在返回 None
        """
        task = await self.get_task(task_id)
        if not task:
            return None

        # 获取所有快照
        result = await self.db.execute(
            select(PerformanceSnapshot)
            .where(PerformanceSnapshot.task_id == task_id)
            .order_by(PerformanceSnapshot.elapsed.asc())
        )
        snapshots = list(result.scalars().all())

        # 汇总状态码分布
        status_2xx = sum(s.status_2xx for s in snapshots)
        status_3xx = sum(s.status_3xx for s in snapshots)
        status_4xx = sum(s.status_4xx for s in snapshots)
        status_5xx = sum(s.status_5xx for s in snapshots)

        return {
            "summary": {
                "task_id": task.id,
                "task_name": task.name,
                "target_url": task.target_url,
                "method": task.method,
                "mode": task.mode,
                "concurrency": task.concurrency,
                "duration_setting": task.duration,
                "duration_actual": task.duration_actual,
                "start_time": task.start_time.isoformat() if task.start_time else None,
                "end_time": task.end_time.isoformat() if task.end_time else None,
                "account_key": task.account_key,
            },
            "metrics": {
                "total_requests": task.total_requests,
                "success_count": task.success_count,
                "fail_count": task.fail_count,
                "error_rate": task.error_rate,
                "actual_tps": task.actual_tps,
                "min_response_time": task.min_response_time,
                "max_response_time": task.max_response_time,
                "avg_response_time": task.avg_response_time,
                "p50_response_time": task.p50_response_time,
                "p90_response_time": task.p90_response_time,
                "p99_response_time": task.p99_response_time,
                "status_2xx": status_2xx,
                "status_3xx": status_3xx,
                "status_4xx": status_4xx,
                "status_5xx": status_5xx,
            },
            "snapshots": [s.to_dict() for s in snapshots],
        }

    # ==================== 内部方法 ====================

    @staticmethod
    def _build_client_config(task_dict: Dict[str, Any]) -> Dict[str, Any]:
        """
        构建 HttpClient 配置

        从任务的目标 URL 推导登录系统的基础地址，
        使用默认的账号密码和认证配置。

        Args:
            task_dict: 任务数据字典

        Returns:
            dict: HttpClient 配置字典
        """
        target_url = task_dict.get("target_url", "")
        # 从 target_url 提取 base_url（scheme + host）
        base_url = _DEFAULT_BASE_URL
        if target_url:
            try:
                from urllib.parse import urlparse
                parsed = urlparse(target_url)
                if parsed.scheme and parsed.netloc:
                    base_url = f"{parsed.scheme}://{parsed.netloc}"
            except Exception:
                pass

        return {
            "target": {
                "base_url": base_url,
                "verify_ssl": False,
                "timeout": 30,
            },
            "accounts": dict(_DEFAULT_ACCOUNTS),
            "auth": dict(_DEFAULT_AUTH_CONFIG),
        }

    @staticmethod
    def _apply_result(task: PerformanceTask, result: Dict[str, Any]) -> None:
        """
        将执行结果写入任务统计字段

        Args:
            task: 任务对象
            result: 执行结果字典
        """
        task.total_requests = result.get("total_requests", 0)
        task.success_count = result.get("success_count", 0)
        task.fail_count = result.get("fail_count", 0)
        task.min_response_time = result.get("min_response_time", 0.0)
        task.max_response_time = result.get("max_response_time", 0.0)
        task.avg_response_time = result.get("avg_response_time", 0.0)
        task.p50_response_time = result.get("p50_response_time", 0.0)
        task.p90_response_time = result.get("p90_response_time", 0.0)
        task.p99_response_time = result.get("p99_response_time", 0.0)
        task.actual_tps = result.get("actual_tps", 0.0)
        task.error_rate = result.get("error_rate", 0.0)

    async def _save_snapshots(
        self, task_id: str, snapshots: List[Dict[str, Any]]
    ) -> int:
        """
        批量保存快照到数据库

        Args:
            task_id: 任务ID
            snapshots: 快照数据列表

        Returns:
            int: 保存的快照数量
        """
        if not snapshots:
            return 0

        batch = []
        for s in snapshots:
            batch.append(PerformanceSnapshot(
                task_id=task_id,
                timestamp=s.get("timestamp", time.time()),
                elapsed=s.get("elapsed", 0.0),
                concurrent_users=s.get("concurrent_users", 0),
                tps=s.get("tps", 0.0),
                avg_response_time=s.get("avg_response_time", 0.0),
                min_response_time=s.get("min_response_time", 0.0),
                max_response_time=s.get("max_response_time", 0.0),
                p50_response_time=s.get("p50", 0.0),
                p90_response_time=s.get("p90", 0.0),
                p99_response_time=s.get("p99", 0.0),
                success_count=s.get("success_count", 0),
                fail_count=s.get("fail_count", 0),
                status_2xx=s.get("status_2xx", 0),
                status_3xx=s.get("status_3xx", 0),
                status_4xx=s.get("status_4xx", 0),
                status_5xx=s.get("status_5xx", 0),
            ))

        self.db.add_all(batch)
        await self.db.flush()
        logger.info(f"批量保存快照: {len(batch)} 条 (task_id={task_id})")
        return len(batch)

    async def _save_outputs(
        self, task_id: str, captured: Optional[Dict[str, List[Any]]]
    ) -> int:
        """
        保存响应捕获输出到数据库

        使用 INSERT ... ON CONFLICT 语义（先查后写），
        每个 task_id+key 唯一，已存在时追加值。

        Args:
            task_id: 任务ID
            captured: 捕获数据 {key: [values]}

        Returns:
            int: 保存/更新的输出条目数
        """
        if not captured:
            return 0

        saved = 0
        for key, values in captured.items():
            if not values:
                continue

            # 查重：同一 task_id+key 已存在则追加
            result = await self.db.execute(
                select(PerformanceTaskOutput)
                .where(
                    PerformanceTaskOutput.task_id == task_id,
                    PerformanceTaskOutput.key == key,
                )
            )
            existing = result.scalar_one_or_none()

            if existing:
                # 合并去重
                existing_values = list(existing.value) if existing.value else []
                seen = set(str(v) for v in existing_values)
                for v in values:
                    key_str = str(v)
                    if key_str not in seen:
                        existing_values.append(v)
                        seen.add(key_str)
                existing.value = existing_values
            else:
                self.db.add(PerformanceTaskOutput(
                    task_id=task_id,
                    key=key,
                    value=list(values),
                ))
            saved += 1

        await self.db.flush()
        logger.info(f"保存响应捕获输出: {saved} 条 (task_id={task_id})")
        return saved

    def _build_ref_provider(self) -> Callable[[str, str], list]:
        """
        构建跨任务引用取值回调

        供 TemplateResolver 的 REF 占位符使用，
        用 asyncio.run_coroutine_threadsafe 在事件循环中查询数据库。

        Returns:
            Callable: (task_id, key) -> list
        """
        db = self.db

        def _ref_provider(task_id: str, key: str) -> list:
            """同步回调，在 executor 工作线程中调用"""
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    future = asyncio.run_coroutine_threadsafe(
                        self._query_output(task_id, key), loop
                    )
                    return future.result(timeout=5) or []
                return []
            except Exception as e:
                logger.warning(f"REF 查询失败 task_id={task_id}, key={key}: {e}")
                return []

        return _ref_provider

    async def _query_output(self, task_id: str, key: str) -> list:
        """
        查询某个任务的输出值

        Args:
            task_id: 任务ID
            key: 输出键名

        Returns:
            list: 输出值列表
        """
        result = await self.db.execute(
            select(PerformanceTaskOutput)
            .where(
                PerformanceTaskOutput.task_id == task_id,
                PerformanceTaskOutput.key == key,
            )
        )
        output = result.scalar_one_or_none()
        return list(output.value) if output and output.value else []