#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：security_service.py
模块描述：安全测试服务层，提供安全测试配置管理、执行调度、结果查询等业务逻辑

作者：czj
创建日期：2026-07-21
最后修改：2026-07-21
"""

import logging
from datetime import datetime
from typing import Optional, List, Tuple, Dict, Any
from concurrent.futures import ThreadPoolExecutor
import asyncio

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

from app.models.security_config import SecurityConfig
from app.models.vulnerability_result import VulnerabilityResult
from app.models.execution import Execution
from app.models.test_case import TestCase
from app.models.module import Module
from app.executors.security_executor import SecurityExecutor, SecurityCaseResult
from app.utils.id_generator import generate_id
from app.config import settings

logger = logging.getLogger(__name__)


class SecurityService:
    """
    安全测试服务类

    提供安全测试配置管理、执行调度、结果查询等功能。
    """

    def __init__(self, db: AsyncSession):
        """
        初始化安全测试服务

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

    # ==================== 配置管理 ====================

    async def get_config(self, config_id: str) -> Optional[SecurityConfig]:
        """
        获取安全测试配置

        Args:
            config_id: 配置ID

        Returns:
            SecurityConfig: 配置对象，不存在返回 None
        """
        result = await self.db.execute(
            select(SecurityConfig).where(SecurityConfig.id == config_id)
        )
        return result.scalar_one_or_none()

    async def get_default_config(self) -> Optional[SecurityConfig]:
        """
        获取默认安全测试配置

        Returns:
            SecurityConfig: 默认配置，不存在返回 None
        """
        result = await self.db.execute(
            select(SecurityConfig).where(SecurityConfig.is_default == True)
        )
        return result.scalar_one_or_none()

    async def list_configs(
        self,
        skip: int = 0,
        limit: int = 20
    ) -> Tuple[List[SecurityConfig], int]:
        """
        获取安全测试配置列表

        Args:
            skip: 跳过记录数
            limit: 返回记录数

        Returns:
            tuple: (配置列表, 总数)
        """
        # 查询总数
        count_result = await self.db.execute(
            select(func.count(SecurityConfig.id))
        )
        total = count_result.scalar() or 0

        # 查询列表
        result = await self.db.execute(
            select(SecurityConfig)
            .order_by(desc(SecurityConfig.is_default), desc(SecurityConfig.updated_at))
            .offset(skip)
            .limit(limit)
        )
        configs = list(result.scalars().all())

        return configs, total

    async def create_config(self, config_data: Dict[str, Any]) -> SecurityConfig:
        """
        创建安全测试配置

        Args:
            config_data: 配置数据

        Returns:
            SecurityConfig: 创建的配置对象
        """
        config_id = generate_id("sec_cfg")

        config = SecurityConfig(
            id=config_id,
            name=config_data.get("name", ""),
            target_url=config_data.get("target_url", ""),
            server_ip=config_data.get("server_ip", ""),
            verify_ssl=config_data.get("verify_ssl", False),
            timeout=config_data.get("timeout", 30),
            accounts=config_data.get("accounts", {}),
            auth_config=config_data.get("auth_config", {}),
            rate_limits=config_data.get("rate_limits", {}),
            api_prefixes=config_data.get("api_prefixes", {}),
            is_default=config_data.get("is_default", False),
            description=config_data.get("description", ""),
        )

        self.db.add(config)
        await self.db.commit()
        await self.db.refresh(config)

        logger.info(f"创建安全测试配置: {config_id} - {config.name}")
        return config

    async def update_config(
        self,
        config_id: str,
        config_data: Dict[str, Any]
    ) -> Optional[SecurityConfig]:
        """
        更新安全测试配置

        Args:
            config_id: 配置ID
            config_data: 更新数据

        Returns:
            SecurityConfig: 更新后的配置，不存在返回 None
        """
        config = await self.get_config(config_id)
        if not config:
            return None

        # 更新字段
        for key, value in config_data.items():
            if value is not None and hasattr(config, key):
                setattr(config, key, value)

        config.updated_at = datetime.utcnow()

        await self.db.commit()
        await self.db.refresh(config)

        logger.info(f"更新安全测试配置: {config_id}")
        return config

    async def delete_config(self, config_id: str) -> bool:
        """
        删除安全测试配置

        Args:
            config_id: 配置ID

        Returns:
            bool: 是否成功删除
        """
        config = await self.get_config(config_id)
        if not config:
            return False

        await self.db.delete(config)
        await self.db.commit()

        logger.info(f"删除安全测试配置: {config_id}")
        return True

    # ==================== 账号测试 ====================

    async def test_account(
        self,
        config_id: str,
        account_key: str
    ) -> Dict[str, Any]:
        """
        测试账号连接

        Args:
            config_id: 配置ID
            account_key: 账号键名

        Returns:
            dict: 测试结果
        """
        config = await self.get_config(config_id)
        if not config:
            return {
                "success": False,
                "account_key": account_key,
                "token_preview": "",
                "error_message": "配置不存在"
            }

        # 创建临时执行器测试登录
        executor = SecurityExecutor(config.to_client_config())
        result = executor.auth.test_account(account_key)

        return {
            "success": result["success"],
            "account_key": account_key,
            "token_preview": result["token_preview"],
            "error_message": result["error"]
        }

    # ==================== 执行管理 ====================

    async def create_execution(
        self,
        config_id: str,
        module_ids: List[str] = None,
        case_ids: List[str] = None,
        name: str = ""
    ) -> Execution:
        """
        创建安全测试执行记录

        Args:
            config_id: 安全测试配置ID
            module_ids: 要执行的模块ID列表
            case_ids: 要执行的用例ID列表
            name: 执行名称

        Returns:
            Execution: 执行记录
        """
        execution_id = generate_id("exec")

        # 如果没有指定用例，查询所有安全测试用例
        if not case_ids and not module_ids:
            case_ids = await self._get_all_security_case_ids()
        elif module_ids and not case_ids:
            case_ids = await self._get_case_ids_by_modules(module_ids)

        total_cases = len(case_ids)

        execution = Execution(
            id=execution_id,
            name=name or f"安全测试-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            trigger_type="manual",
            trigger_by="security_test",
            total_cases=total_cases,
            status="pending",
            case_type="security",
            config={"config_id": config_id, "case_ids": case_ids},
        )

        self.db.add(execution)
        await self.db.commit()
        await self.db.refresh(execution)

        logger.info(f"创建安全测试执行: {execution_id}, 用例数: {total_cases}")
        return execution

    async def _get_all_security_case_ids(self) -> List[str]:
        """获取所有安全测试用例ID"""
        result = await self.db.execute(
            select(TestCase.id).where(TestCase.case_type == "security")
        )
        return [row[0] for row in result.all()]

    async def _get_case_ids_by_modules(self, module_ids: List[str]) -> List[str]:
        """获取指定模块的安全测试用例ID"""
        result = await self.db.execute(
            select(TestCase.id)
            .where(TestCase.module_id.in_(module_ids))
            .where(TestCase.case_type == "security")
        )
        return [row[0] for row in result.all()]

    async def run_execution(self, execution_id: str, progress_callback=None):
        """
        执行安全测试

        Args:
            execution_id: 执行ID
            progress_callback: 进度回调函数
        """
        execution = await self._get_execution(execution_id)
        if not execution:
            logger.error(f"执行记录不存在: {execution_id}")
            return

        config_id = execution.config.get("config_id")
        case_ids = execution.config.get("case_ids", [])

        # 获取配置
        config = await self.get_config(config_id)
        if not config:
            logger.error(f"安全测试配置不存在: {config_id}")
            execution.status = "failed"
            await self.db.commit()
            return

        # 获取用例
        cases = await self._get_cases_by_ids(case_ids)
        if not cases:
            logger.warning(f"没有找到要执行的用例")
            execution.status = "completed"
            await self.db.commit()
            return

        # 更新执行状态
        execution.status = "running"
        execution.start_time = datetime.utcnow()
        await self.db.commit()

        try:
            # 在线程池中执行（避免阻塞 asyncio）
            loop = asyncio.get_event_loop()

            # 在 async 上下文中预先把 ORM 对象转成纯字典，避免：
            # 1. SQLAlchemy JSON 字段在离开 session 上下文后变为 str
            # 2. ORM 对象在 run_in_executor 中脱钩（detached）引发异常
            import json as _json

            def _to_plain_steps(steps_val):
                if steps_val is None:
                    return {}
                if isinstance(steps_val, str):
                    try:
                        parsed = _json.loads(steps_val)
                        return parsed if isinstance(parsed, dict) else {}
                    except Exception:
                        return {}
                if isinstance(steps_val, dict):
                    return steps_val
                return {}

            plain_cases = [
                {
                    "id": c.id,
                    "name": c.name,
                    "steps": _to_plain_steps(c.steps),
                }
                for c in cases
            ]
            client_config = config.to_client_config()

            # 同步执行函数
            def sync_run():
                try:
                    executor = SecurityExecutor(client_config)
                    executor.start()
                    return executor.execute_cases(
                        plain_cases,
                        progress_callback=lambda cid, st, r: None  # 同步回调先不处理
                    )
                except Exception as e:
                    import traceback
                    logger.error(f"sync_run 异常: {e}")
                    logger.error(traceback.format_exc())
                    raise

            # 在线程池中执行
            results: List[SecurityCaseResult] = await loop.run_in_executor(
                None, sync_run
            )

            # 保存结果
            passed = 0
            failed = 0
            for r in results:
                await self._save_vulnerability_result(execution_id, r)
                if r.is_vulnerable:
                    failed += 1
                else:
                    passed += 1

            # 更新执行统计
            execution.passed = passed
            execution.failed = failed
            execution.status = "completed"
            execution.end_time = datetime.utcnow()
            execution.duration = (execution.end_time - execution.start_time).total_seconds()
            execution.pass_rate = round((passed / execution.total_cases) * 100, 2) if execution.total_cases > 0 else 0

            await self.db.commit()
            logger.info(f"安全测试执行完成: {execution_id}, 通过: {passed}, 失败: {failed}")

            # 自动生成 Markdown 报告（失败不影响执行结果）
            try:
                report_path = await self._generate_report(execution_id)
                if report_path:
                    # 将报告路径写入 execution.config
                    await self.db.refresh(execution)
                    exec_config = dict(execution.config or {})
                    exec_config["report_path"] = report_path
                    execution.config = exec_config
                    await self.db.commit()
                    logger.info(f"安全测试报告已生成: {report_path}")
            except Exception as report_err:
                logger.warning(f"生成安全测试报告失败（不影响执行结果）: {report_err}")

        except Exception as e:
            logger.error(f"安全测试执行异常: {e}")
            execution.status = "failed"
            execution.end_time = datetime.utcnow()
            await self.db.commit()

    async def _generate_report(self, execution_id: str) -> Optional[str]:
        """生成安全测试 Markdown 报告"""
        from app.services.security_report_service import SecurityReportService
        report_service = SecurityReportService()
        return await report_service.generate_report(execution_id, self.db)

    async def _save_vulnerability_result(
        self,
        execution_id: str,
        result: SecurityCaseResult
    ):
        """保存漏洞测试结果"""
        vuln_result = VulnerabilityResult(
            id=generate_id("vuln"),
            execution_id=execution_id,
            case_id=result.case_id,
            test_id=result.test_id,
            name=result.name,
            level=result.level,
            description=result.description,
            request_info=result.request_info,
            response_info=result.response_info,
            is_vulnerable=result.is_vulnerable,
            fix_suggestion=result.fix_suggestion,
            duration=result.duration,
            status=result.status,
            error_message=result.error,
            extra_data=result.metadata,
        )

        self.db.add(vuln_result)
        # 不在这里 commit，由调用方统一处理

    async def _get_execution(self, execution_id: str) -> Optional[Execution]:
        """获取执行记录"""
        result = await self.db.execute(
            select(Execution).where(Execution.id == execution_id)
        )
        return result.scalar_one_or_none()

    async def _get_cases_by_ids(self, case_ids: List[str]) -> List[TestCase]:
        """获取指定ID的用例"""
        if not case_ids:
            return []

        result = await self.db.execute(
            select(TestCase).where(TestCase.id.in_(case_ids))
        )
        return list(result.scalars().all())

    # ==================== 结果查询 ====================

    async def get_vulnerability_results(
        self,
        execution_id: str,
        skip: int = 0,
        limit: int = 50,
        level: Optional[str] = None,
        is_vulnerable: Optional[bool] = None
    ) -> Tuple[List[VulnerabilityResult], int]:
        """
        获取漏洞测试结果列表

        Args:
            execution_id: 执行ID
            skip: 跳过记录数
            limit: 返回记录数
            level: 风险等级筛选
            is_vulnerable: 是否存在漏洞筛选

        Returns:
            tuple: (结果列表, 总数)
        """
        # 构建查询条件
        conditions = [VulnerabilityResult.execution_id == execution_id]

        if level:
            conditions.append(VulnerabilityResult.level == level)
        if is_vulnerable is not None:
            conditions.append(VulnerabilityResult.is_vulnerable == is_vulnerable)

        # 查询总数
        count_query = select(func.count(VulnerabilityResult.id)).where(*conditions)
        count_result = await self.db.execute(count_query)
        total = count_result.scalar() or 0

        # 查询列表（按风险等级排序）
        result = await self.db.execute(
            select(VulnerabilityResult)
            .where(*conditions)
            .order_by(VulnerabilityResult.level, VulnerabilityResult.test_id)
            .offset(skip)
            .limit(limit)
        )
        results = list(result.scalars().all())

        return results, total

    async def get_execution_summary(self, execution_id: str) -> Dict[str, Any]:
        """
        获取执行摘要

        Args:
            execution_id: 执行ID

        Returns:
            dict: 摘要数据
        """
        # 统计各风险等级数量
        result = await self.db.execute(
            select(
                VulnerabilityResult.level,
                func.count(VulnerabilityResult.id),
                func.sum(
                    func.case((VulnerabilityResult.is_vulnerable == True, 1), else_=0)
                )
            )
            .where(VulnerabilityResult.execution_id == execution_id)
            .group_by(VulnerabilityResult.level)
        )

        level_stats = {}
        for row in result.all():
            level, total, vuln_count = row
            level_stats[level] = {
                "total": total,
                "vulnerable": vuln_count or 0,
            }

        return {
            "execution_id": execution_id,
            "level_stats": level_stats,
        }
