import sys
import os
# 将上一级目录加入模块搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import asyncio
import aiohttp
import time
import traceback
import datetime
import pandas as pd
import pymysql
import random
from tqdm import tqdm
from Util.dingtalk_helper import DingTalkHelper

# 将上一级目录加入模块搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# --- 配置 ---
ACCESS_TOKEN = '4625f6690acd9347fae5b3a05af598be63e73d604b933a9b3902425b8f136d4d'
SECRET = 'SEC3b6937550bd297b5491855f6f40c2ff1b41bc8c495e118ba9848742b1ddf8f19'

apiname = "pda笼位更新"
url = "http://qilu.baoyizn.com:9935/api/base/cage/cage/trainingCage"
headers = {
    "user-agent": "Mozilla/5.0 (Linux; Android 10; DS-MDT301 Build/QP1A.190711.020; wv) AppleWebKit/537.36 (KHTML, like Gecko) Version/4.0 Chrome/74.0.3729.186 Mobile Safari/537.36 uni-app Html5Plus/1.0 (Immersed/24.0)",
    "Content-Type": "application/json; charset=utf-8",
    "Accept-Encoding": "gzip"
}

NUM_WORKERS = 100
TOTAL_REQUESTS = 20000
MAX_RETRIES = 3
REQUEST_TIMEOUT = 60
OUTPUT_DIR = './load_test_report'

# --- 数据库配置 ---
# DB_CONFIG = {
#     'host': 'rm-wz97y18i4t16hfsk6qo.mysql.rds.aliyuncs.com',
#     'database': 'srps_qilu',
#     'user': 'dev',
#     'password': 'Hello@112',
#     'charset': 'utf8mb4'
# }
DB_CONFIG = {
    'host': 'rm-wz90i533p18631e685o.mysql.rds.aliyuncs.com',
    'database': 'srps_qilu_prod',
    'user': 'dev',
    'password': 'Hello@112',
    'charset': 'utf8mb4'
}

# --- 初始化 ---
dingtalk_helper = DingTalkHelper(ACCESS_TOKEN, SECRET)

class DataManager:
    """数据管理器，负责从数据库和Excel加载数据"""

    def __init__(self):
        self.cage_ids = []
        self.user_data = []  # 存储用户ID和对应的课题组ID
        self.user_tokens = []  # 存储用户token

    def load_user_tokens(self, excel_file='TestData/userinfo.xlsx'):
        """从Excel加载用户token"""
        try:
            df = pd.read_excel(excel_file)
            self.user_tokens = df['token'].tolist()
            print(f"成功加载 {len(self.user_tokens)} 个用户token")
            return True
        except Exception as e:
            print(f"加载用户token失败: {e}")
            return False

    def load_database_data(self):
        """从数据库加载笼位ID和用户信息"""
        try:
            conn = pymysql.connect(**DB_CONFIG)

            # 获取笼位ID
            with conn.cursor() as cursor:
                cursor.execute("SELECT id FROM l_cage WHERE id IS NOT NULL")
                self.cage_ids = [row[0] for row in cursor.fetchall()]
                print(f"成功加载 {len(self.cage_ids)} 个笼位ID")

            # 获取用户ID和课题组ID
            with conn.cursor() as cursor:
                cursor.execute("SELECT id, research_group_ids FROM sys_user WHERE research_group_ids IS NOT NULL AND research_group_ids != ''")

                for user_id, research_groups in cursor.fetchall():
                    # 处理多个课题组ID（用逗号分隔的情况）
                    if research_groups and ',' in research_groups:
                        group_ids = [gid.strip() for gid in research_groups.split(',') if gid.strip()]
                    elif research_groups:
                        group_ids = [research_groups.strip()]
                    else:
                        continue

                    self.user_data.append({
                        'user_id': user_id,
                        'research_group_ids': group_ids
                    })

                print(f"成功加载 {len(self.user_data)} 个有效用户数据")

            conn.close()
            return len(self.cage_ids) > 0 and len(self.user_data) > 0

        except Exception as e:
            print(f"数据库连接失败: {e}")
            return False

    def get_random_cage_id(self):
        """随机获取一个笼位ID"""
        return random.choice(self.cage_ids) if self.cage_ids else None

    def get_random_user_data(self):
        """随机获取一个用户数据（用户ID和课题组ID）"""
        if not self.user_data:
            return None, None

        user = random.choice(self.user_data)
        user_id = user['user_id']
        research_group_id = random.choice(user['research_group_ids'])

        return user_id, research_group_id

    def get_random_token(self):
        """随机获取一个用户token"""
        return random.choice(self.user_tokens) if self.user_tokens else None

# 全局数据管理器
data_manager = DataManager()

def create_cage_data():
    """创建笼位更新请求数据"""
    cage_id = data_manager.get_random_cage_id()
    user_id, research_group_id = data_manager.get_random_user_data()

    if not cage_id or not user_id or not research_group_id:
        return None

    return {
        "cage": {"id": str(cage_id)},
        "researchGroup": {"id": research_group_id},
        "user": {"id": user_id}
    }

async def perform_request(session: aiohttp.ClientSession, index: int, max_retries: int = MAX_RETRIES):
    """执行单个请求"""
    attempt = 0
    last_err = None
    token = data_manager.get_random_token()

    if not token:
        return {
            'index': index,
            'timestamp': time.time(),
            'status_code': 0,
            'latency_ms': 0,
            'response_size': None,
            'error': 'No available token'
        }

    # 动态设置token到headers
    current_headers = headers.copy()
    current_headers['token'] = token

    while attempt < max_retries:
        data = create_cage_data()
        if not data:
            return {
                'index': index,
                'timestamp': time.time(),
                'status_code': 0,
                'latency_ms': 0,
                'response_size': None,
                'error': 'No available cage/user data'
            }

        start = time.time()
        try:
            async with session.post(url, json=data, headers=current_headers) as resp:
                text = await resp.text()
                latency_ms = (time.time() - start) * 1000.0
                status = resp.status
                if status == 200:
                    return {
                        'index': index,
                        'timestamp': time.time(),
                        'status_code': status,
                        'latency_ms': latency_ms,
                        'response_size': len(text) if text is not None else None,
                        'error': None
                    }
                else:
                    last_err = f'status_{status}:{text}'
                    attempt += 1
                    await asyncio.sleep(min(10, 2 ** attempt))
        except Exception as e:
            latency_ms = (time.time() - start) * 1000.0
            last_err = f'{type(e).__name__}:{str(e)}'
            attempt += 1
            await asyncio.sleep(min(10, 2 ** attempt))

    # 最终失败
    return {
        'index': index,
        'timestamp': time.time(),
        'status_code': 0,
        'latency_ms': latency_ms if 'latency_ms' in locals() else 0,
        'response_size': None,
        'error': last_err
    }

async def worker(name: int, queue: asyncio.Queue, session: aiohttp.ClientSession, gen, pbar, success_counter: dict, failed_list: list, lock: asyncio.Lock):
    """工作线程"""
    while True:
        idx = await queue.get()
        if idx is None:
            queue.task_done()
            break
        try:
            res = await perform_request(session, idx)
            # 记录到报告生成器
            if gen:
                gen.record_result(
                    index=res['index'],
                    timestamp=res['timestamp'],
                    status_code=int(res['status_code']),
                    latency_ms=float(res['latency_ms']),
                    response_size=res.get('response_size'),
                    error=res.get('error')
                )
            async with lock:
                if res['status_code'] and 200 <= res['status_code'] < 300:
                    success_counter['count'] += 1
                else:
                    failed_list.append((res['index'], res.get('error')))
                pbar.update(1)
        except Exception as e:
            async with lock:
                failed_list.append((idx, f'Worker异常:{type(e).__name__}:{e}'))
                pbar.update(1)
        finally:
            queue.task_done()

async def batch_update_cages(total: int, num_workers: int):
    """批量更新笼位"""
    # 动态加载报告生成器模块
    gen = None
    try:
        import importlib.util

        # 使用您提供的绝对路径
        report_path = 'H:\\项目\\造数脚本\\Util\\stress_test_report_generator.py'

        if os.path.exists(report_path):
            spec = importlib.util.spec_from_file_location('report_module', report_path)
            report_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(report_module)
            LoadTestReportGenerator = getattr(report_module, 'LoadTestReportGenerator')
            gen = LoadTestReportGenerator(test_name=f'{apiname}压测任务', report_title='压测详细报告')
            print(f"成功加载压测报告生成器: {report_path}")
        else:
            # 备用：尝试相对路径
            script_dir = os.path.dirname(os.path.abspath(__file__))
            report_path = os.path.join(script_dir, 'Util', 'stress_test_report_generator.py')
            if os.path.exists(report_path):
                spec = importlib.util.spec_from_file_location('report_module', report_path)
                report_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(report_module)
                LoadTestReportGenerator = getattr(report_module, 'LoadTestReportGenerator')
                gen = LoadTestReportGenerator(test_name=f'{apiname}压测任务', report_title='压测详细报告')
                print(f"成功加载压测报告生成器: {report_path}")
            else:
                print('警告: 未找到压测报告生成器，将跳过报告生成')
    except Exception as e:
        print(f'加载压测报告生成器失败: {e}，将跳过报告生成')

    # 初始化数据
    if not data_manager.load_user_tokens():
        print("加载用户token失败，退出压测")
        return

    if not data_manager.load_database_data():
        print("加载数据库数据失败，退出压测")
        return

    timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)
    connector = aiohttp.TCPConnector(limit=num_workers, limit_per_host=num_workers, force_close=False)

    async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
        queue = asyncio.Queue()
        for i in range(1, total + 1):
            await queue.put(i)
        for _ in range(num_workers):
            await queue.put(None)

        success_counter = {'count': 0}
        failed_list = []
        lock = asyncio.Lock()

        with tqdm(total=total, desc='压测进度') as pbar:
            workers = [
                asyncio.create_task(worker(i, queue, session, gen, pbar, success_counter, failed_list, lock))
                for i in range(num_workers)
            ]
            await asyncio.gather(*workers)

        # 任务完成，生成报告
        if gen:
            os.makedirs(OUTPUT_DIR, exist_ok=True)
            outputs = gen.generate_report(OUTPUT_DIR, formats=['html', 'json', 'csv', 'docx'])
            stats = gen.compute_stats()

            # 构造钉钉摘要消息
            now_str = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            msg = [f'【{apiname} 压测报告】', f'生成时间：{now_str}']
            msg.append(f"总请求数：{stats.get('total_requests',0)}，成功：{stats.get('success_count',0)}，失败：{stats.get('fail_count',0)}，成功率：{stats.get('success_rate',0):.2%}")

            if 'duration_seconds' in stats:
                msg.append(f"总耗时(s)：{stats.get('duration_seconds',0):.2f}，平均吞吐(req/s)：{stats.get('throughput_rps',0):.2f}")

            if 'latency_ms' in stats:
                lat = stats.get('latency_ms', {})
                msg.append(f"延迟(ms) - 平均：{lat.get('avg',0):.2f}，P90：{lat.get('p90',0):.2f}，P95：{lat.get('p95',0):.2f}，P99：{lat.get('p99',0):.2f}")

            msg.append(f"使用用户数：{len(data_manager.user_tokens)}")
            msg.append(f"可用笼位数：{len(data_manager.cage_ids)}")
            msg.append(f"有效用户数：{len(data_manager.user_data)}")

            # 列出生成的报告文件
            if outputs:
                file_list = []
                for k, v in outputs.items():
                    if k == 'charts':
                        for cname, cpath in v.items():
                            file_list.append(os.path.abspath(cpath))
                    else:
                        file_list.append(os.path.abspath(v))
                msg.append('生成文件：')
                for p in file_list:
                    msg.append(p)

            final_msg = '\n'.join(msg)
        else:
            # 如果没有报告生成器，计算基础统计
            stats = {
                'total_requests': total,
                'success_count': success_counter['count'],
                'fail_count': len(failed_list),
                'success_rate': success_counter['count'] / total if total > 0 else 0
            }

            # 基础统计消息
            now_str = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            msg = [f'【{apiname} 压测报告】', f'生成时间：{now_str}']
            msg.append(f"总请求数：{stats.get('total_requests',0)}，成功：{stats.get('success_count',0)}，失败：{stats.get('fail_count',0)}，成功率：{stats.get('success_rate',0):.2%}")
            msg.append(f"使用用户数：{len(data_manager.user_tokens)}")
            msg.append(f"可用笼位数：{len(data_manager.cage_ids)}")
            msg.append(f"有效用户数：{len(data_manager.user_data)}")
            final_msg = '\n'.join(msg)

        # 发送钉钉消息
        try:
            dingtalk_helper.send_message(final_msg)
            print('钉钉消息发送成功')
        except Exception as e:
            print('发送钉钉消息失败：', e)

        print('\n[压测摘要]')
        print('成功数：', success_counter['count'], ' 失败数：', len(failed_list))
        if failed_list:
            print('失败示例（最多显示20条）：')
            for idx, err in failed_list[:20]:
                print(f'  #{idx} => {err}')

if __name__ == '__main__':
    # 运行前检查依赖
    try:
        import pymysql
        import pandas
        print("依赖检查通过，开始压测...")
    except ImportError as e:
        print(f"缺少依赖: {e}")
        print("请安装: pip install pymysql pandas")
        exit(1)

    # 运行压测
    asyncio.run(batch_update_cages(TOTAL_REQUESTS, NUM_WORKERS))