"""
并发测试：批量创建浏览器环境并执行广告点击+聊天操作
"""

from loguru import logger
from adspower_client import AdsPowerClient
from config import Config
from db_manager import SiteManager, ClickManager, InteractionManager
import sys
import time
import random
from datetime import datetime
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict
import threading  # 添加线程锁支持

# 配置日志
logger.remove()
logger.add(
    sys.stdout,
    format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
    level="INFO"
)


class ConcurrentTester:
    """并发测试管理器"""
    
    # 类级别的浏览器启动锁，确保启动操作串行化
    _browser_start_lock = threading.Lock()
    
    def __init__(self, test_url: str, max_workers: int = 3):
        """
        初始化并发测试器
        
        Args:
            test_url: 测试的目标URL
            max_workers: 最大并发数
        """
        self.test_url = test_url
        self.max_workers = max_workers
        self.client = AdsPowerClient()
        self.created_profiles = []  # 记录创建的环境ID
        self.created_proxies = []  # 记录创建的代理ID
        
        # 创建测试目录
        self.test_base_dir = Path("./test_concurrent")
        self.test_base_dir.mkdir(exist_ok=True)
        
        # 初始化数据库
        self.site_mgr = SiteManager()
        self.site_id = self._init_site()
        
    def _init_site(self) -> int:
        """初始化或获取站点"""
        site = self.site_mgr.get_site_by_url(self.test_url)
        if not site:
            site_id = self.site_mgr.add_site(
                site_url=self.test_url,
                site_name="并发测试站点",
                site_dimension="医疗健康"
            )
            logger.info(f"✅ 创建测试站点: site_id={site_id}")
        else:
            site_id = site['id']
            logger.info(f"✅ 使用已存在站点: site_id={site_id}")
        return site_id
    
    def create_browser_profile(self, index: int) -> Dict:
        """
        创建浏览器环境
        
        Args:
            index: 环境编号
            
        Returns:
            环境信息字典
        """
        try:
            # 获取分组ID
            group_id = self.client.get_group_by_env()
            time.sleep(0.5)  # API 调用间隔
            
            # 获取大麦IP代理
            logger.info(f"[环境 {index}] 获取代理IP...")
            proxy_info = self.client.get_damai_proxy()
            time.sleep(0.5)  # API 调用间隔
            
            if not proxy_info:
                logger.warning(f"[环境 {index}] 获取代理失败，将使用随机代理")
                proxy_config = {}
                proxy_id = None
            else:
                logger.info(f"[环境 {index}] 代理IP: {proxy_info['host']}:{proxy_info['port']}")
                
                # 创建代理并记录ID
                proxy_data = {
                    "type": "http",
                    "host": proxy_info["host"],
                    "port": proxy_info["port"],
                    "user": self.client.DAMAI_USER,
                    "password": self.client.DAMAI_PASSWORD,
                    "remark": f"并发测试代理_{index}"
                }
                
                proxy_id = self.client.create_proxy(proxy_data)
                time.sleep(0.5)  # API 调用间隔
                
                if proxy_id:
                    self.created_proxies.append(proxy_id)
                    logger.info(f"[环境 {index}] 创建代理: {proxy_id}")
                    proxy_config = {"proxyid": proxy_id}
                else:
                    logger.warning(f"[环境 {index}] 创建代理失败")
                    proxy_config = {}
                    proxy_id = None
            
            profile_data = {
                "name": f"并发测试_{index}_{datetime.now().strftime('%H%M%S')}",
                "group_id": str(group_id) if group_id else "0",
                "platform": "health.baidu.com",
                "tabs": [self.test_url],
                "repeat_config": [],
                "ignore_cookie_error": "1",
                "country": "cn",
                "city": "beijing",
                "remark": f"并发测试环境 #{index}",
                "fingerprint_config": {
                    "automatic_timezone": "1",
                    "flash": "block",
                    "scan_port_type": "1",
                    "location": "ask",
                    "location_switch": "1",
                    "canvas": "0",
                    "webgl": "0",
                    "audio": "0",
                    "webrtc": "local",
                    "do_not_track": "true",
                    "hardware_concurrency": "default",
                    "device_memory": "default",
                    "gpu": "2",
                    "mac_address_config": {
                        "model": "1",
                        "address": ""
                    },
                    "browser_kernel_config": {
                        "version": "latest",
                        "type": "chrome"
                    },
                    "random_ua": {
                        "ua_system_version": ["Windows"]
                    }
                }
            }
            
            # 如果有代理配置，添加到profile_data
            if proxy_config:
                profile_data.update(proxy_config)
            
            response = self.client._make_request(
                'POST',
                '/api/v2/browser-profile/create',
                json=profile_data
            )
            
            if response and response.get('code') == 0:
                profile_id = response.get('data', {}).get('profile_id')
                logger.info(f"✅ 创建环境 #{index}: {profile_id}")
                self.created_profiles.append(profile_id)
                return {
                    'index': index,
                    'profile_id': profile_id,
                    'name': profile_data['name'],
                    'proxy': proxy_info if proxy_info else None
                }
            else:
                logger.error(f"❌ 创建环境 #{index} 失败: {response}")
                return None
                
        except Exception as e:
            logger.error(f"❌ 创建环境 #{index} 异常: {str(e)}")
            return None
    
    def run_single_task(self, profile_info: Dict) -> Dict:
        """
        执行单个浏览器任务
        
        Args:
            profile_info: 环境信息
            
        Returns:
            执行结果
        """
        index = profile_info['index']
        profile_id = profile_info['profile_id']
        
        result = {
            'index': index,
            'profile_id': profile_id,
            'success': False,
            'click_id': None,
            'interaction_id': None,
            'error': None
        }
        
        # 创建任务文件夹
        task_folder = self.test_base_dir / f"task_{index}_{datetime.now().strftime('%H%M%S')}"
        task_folder.mkdir(exist_ok=True)
        
        # 每个线程创建自己的 AdsPowerClient 实例
        client = AdsPowerClient()
        
        try:
            logger.info(f"[任务 {index}] 启动浏览器: {profile_id}")
            
            # 使用类锁确保浏览器启动串行化，避免 API 频率限制
            with self._browser_start_lock:
                logger.debug(f"[任务 {index}] 获取启动锁...")
                # 启动浏览器
                browser_info = client.start_browser(user_id=profile_id)
                if not browser_info:
                    result['error'] = "启动浏览器失败"
                    return result
                # 启动后等待，避免下一个启动请求过快
                time.sleep(1.5)
                logger.debug(f"[任务 {index}] 释放启动锁")
            
            time.sleep(1)  # 额外等待浏览器完全启动
            
            # 连接浏览器
            browser = client.connect_browser(browser_info)
            if not browser:
                result['error'] = "CDP连接失败"
                return result
            
            # 获取页面
            context = browser.contexts[0]
            pages = context.pages
            
            # 清理多余页面
            for p in pages:
                if 'start.adspower.net' in p.url:
                    pages.remove(p)
            
            if pages:
                page = pages[0]
            else:
                page = context.new_page()
            
            logger.info(f"[任务 {index}] 访问页面: {self.test_url}")
            page.goto(self.test_url, wait_until='domcontentloaded', timeout=60000)
            time.sleep(3)
            
            # 等待页面完全加载
            try:
                page.wait_for_load_state('networkidle', timeout=10000)
            except Exception:
                logger.warning(f"[任务 {index}] 网络空闲超时，继续执行")
                time.sleep(2)
            
            # 截图
            page.screenshot(path=str(task_folder / "01_loaded.png"))
            
            # 查找并点击广告
            ad_selector = 'span.ec-tuiguang.ecfc-tuiguang.xz81bbe'
            ad_elements = page.locator(ad_selector)
            ad_count = ad_elements.count()
            
            logger.info(f"[任务 {index}] 找到 {ad_count} 个广告")
            
            if ad_count > 0:
                # 点击第一个广告
                first_ad = ad_elements.first
                first_ad.scroll_into_view_if_needed()
                time.sleep(1)
                
                # 点击，忽略超时错误
                try:
                    first_ad.click(timeout=60000)
                    logger.info(f"[任务 {index}] ✅ 已点击广告")
                except Exception as click_err:
                    logger.warning(f"[任务 {index}] 点击超时，但可能已跳转")
                
                # 记录点击
                click_mgr = ClickManager()
                click_id = click_mgr.record_click(
                    site_id=self.site_id,
                    site_url=self.test_url,
                    user_ip=None,
                    device_type='pc'
                )
                result['click_id'] = click_id
                
                # 等待跳转
                time.sleep(3)
                page.wait_for_load_state('domcontentloaded')
                page.screenshot(path=str(task_folder / "02_after_click.png"))
                
                # 发送消息
                messages = [
                    "我想要预约一个医生，有什么推荐吗？",
                    "我现在本人不在当地，医生什么时候有空，是随时能去吗？",
                    "咱们医院是周六日是否上班，随时去吗？",
                    "想找医生看看，有没有推荐的医生",
                    "最近很不舒服，也说不出来全部的症状，能不能直接对话医生？"
                ]
                message = random.choice(messages)
                
                # 滚动到底部
                page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
                time.sleep(1)
                
                # 查找输入框
                input_selectors = [
                    "textarea[contenteditable='true']",
                    "textarea",
                    "textarea[placeholder]",
                    "input[type='text']"
                ]
                
                input_found = False
                for selector in input_selectors:
                    try:
                        count = page.locator(selector).count()
                        if count > 0:
                            for i in range(count):
                                input_elem = page.locator(selector).nth(i)
                                if input_elem.is_visible(timeout=1000):
                                    input_elem.scroll_into_view_if_needed()
                                    time.sleep(0.5)
                                    input_elem.click()
                                    time.sleep(0.5)
                                    input_elem.fill(message)
                                    logger.info(f"[任务 {index}] ✅ 已输入消息")
                                    input_found = True
                                    break
                        if input_found:
                            break
                    except:
                        continue
                
                # 兜底方案
                if not input_found:
                    logger.warning(f"[任务 {index}] 未找到输入框，尝试兜底方案...")
                    try:
                        # 检查 viewport_size 是否为 None
                        if page.viewport_size is None:
                            logger.warning(f"[任务 {index}] viewport_size 为 None，设置默认视口")
                            # 设置默认视口大小
                            page.set_viewport_size({"width": 1280, "height": 720})
                            time.sleep(0.5)
                        
                        viewport_height = page.viewport_size['height']
                        click_x = page.viewport_size['width'] // 2
                        click_y = viewport_height - 10
                        
                        logger.debug(f"[任务 {index}] 点击位置: ({click_x}, {click_y})")
                        page.mouse.click(click_x, click_y)
                        time.sleep(1)
                        page.keyboard.type(message, delay=50)
                        logger.info(f"[任务 {index}] ✅ 已输入消息（兜底）")
                        input_found = True
                    except Exception as fallback_err:
                        logger.error(f"[任务 {index}] 兜底方案失败: {str(fallback_err)}")
                
                # 发送消息
                if input_found:
                    try:
                        page.keyboard.press('Enter')
                        logger.info(f"[任务 {index}] ✅ 已发送消息")
                        time.sleep(2)
                        
                        # 记录互动
                        interaction_mgr = InteractionManager()
                        interaction_id = interaction_mgr.record_interaction(
                            site_id=self.site_id,
                            click_id=click_id,
                            interaction_type='message',
                            reply_content=message,
                            is_successful=True,
                            response_received=False,
                            response_content=None
                        )
                        result['interaction_id'] = interaction_id
                        
                        page.screenshot(path=str(task_folder / "03_sent.png"))
                        result['success'] = True
                    except Exception as e:
                        logger.warning(f"[任务 {index}] 发送失败: {str(e)}")
            
            # 关闭浏览器前，截图聊天页面最终状态
            try:
                logger.info(f"[任务 {index}] 截图聊天页面...")
                # 等待可能的回复消息加载
                time.sleep(2)
                # 滚动到页面顶部，确保看到完整对话
                page.evaluate("window.scrollTo(0, 0)")
                time.sleep(0.5)
                # 截图整个页面
                page.screenshot(path=str(task_folder / "04_final_chat.png"), full_page=True)
                logger.info(f"[任务 {index}] ✅ 聊天页面截图已保存")
            except Exception as screenshot_err:
                logger.warning(f"[任务 {index}] 截图失败: {str(screenshot_err)}")
            
            # 优雅关闭 Playwright 连接，避免 CancelledError
            try:
                if browser:
                    logger.debug(f"[任务 {index}] 关闭 Playwright 浏览器连接...")
                    browser.close()
                    time.sleep(0.5)
            except Exception as close_err:
                logger.debug(f"[任务 {index}] 关闭浏览器连接异常: {str(close_err)}")
            
            # 根据配置决定是否关闭浏览器进程
            if Config.AUTO_CLOSE_BROWSER:
                try:
                    client.stop_browser(user_id=profile_id)
                    logger.info(f"[任务 {index}] 浏览器已关闭")
                except Exception as stop_err:
                    logger.warning(f"[任务 {index}] 关闭浏览器失败: {str(stop_err)}")
            
        except Exception as e:
            logger.error(f"[任务 {index}] 执行异常: {str(e)}")
            result['error'] = str(e)
            import traceback
            traceback.print_exc()
        
        return result
    
    def delete_profiles(self, profile_ids: List[str]):
        """
        批量删除环境
        
        Args:
            profile_ids: 环境ID列表
        """
        if not profile_ids:
            return
        
        try:
            response = self.client._make_request(
                'POST',
                '/api/v2/browser-profile/delete',
                json={'profile_id': profile_ids}
            )
            
            if response and response.get('code') == 0:
                logger.info(f"✅ 已删除 {len(profile_ids)} 个环境")
            else:
                logger.error(f"❌ 删除环境失败: {response}")
        except Exception as e:
            logger.error(f"❌ 删除环境异常: {str(e)}")
    
    def delete_proxies(self, proxy_ids: List[str]):
        """
        批量删除代理
        
        Args:
            proxy_ids: 代理ID列表
        """
        if not proxy_ids:
            return
        
        try:
            response = self.client._make_request(
                'POST',
                '/api/v2/proxy-list/delete',
                json={'proxy_id': proxy_ids}
            )
            
            if response and response.get('code') == 0:
                logger.info(f"✅ 已删除 {len(proxy_ids)} 个代理")
            else:
                logger.error(f"❌ 删除代理失败: {response}")
        except Exception as e:
            logger.error(f"❌ 删除代理异帰常: {str(e)}")
    
    def run_concurrent_test(self, num_tasks: int):
        """
        运行并发测试
        
        Args:
            num_tasks: 并发任务数量
        """
        logger.info("=" * 60)
        logger.info(f"开始并发测试: {num_tasks} 个任务, 最大并发数: {self.max_workers}")
        logger.info("=" * 60)
        
        # 第一步：批量创建环境
        logger.info("\n步骤 1: 创建浏览器环境")
        logger.info("-" * 60)
        
        profiles = []
        for i in range(num_tasks):
            profile_info = self.create_browser_profile(i + 1)
            if profile_info:
                profiles.append(profile_info)
                proxy_info = profile_info.get('proxy')
                if proxy_info:
                    logger.info(f"环境 #{i+1} 使用代理: {proxy_info['host']}:{proxy_info['port']}")
            # 增加环境创建间隔，避免触发 API 频率限制
            time.sleep(3)  # 每个环境创建后等待 3 秒
        
        logger.info(f"✅ 成功创建 {len(profiles)} 个环境\n")
        
        if not profiles:
            logger.error("没有成功创建任何环境，退出测试")
            return
        
        # 第二步：并发执行任务
        logger.info("步骤 2: 并发执行任务")
        logger.info("-" * 60)
        
        results = []
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = {
                executor.submit(self.run_single_task, profile): profile
                for profile in profiles
            }
            
            for future in as_completed(futures):
                profile = futures[future]
                try:
                    result = future.result()
                    results.append(result)
                    
                    status = "✅ 成功" if result['success'] else "❌ 失败"
                    logger.info(f"[任务 {result['index']}] {status}")
                    
                except Exception as e:
                    logger.error(f"[任务 {profile['index']}] 执行异常: {str(e)}")
        
        # 第三步：统计结果
        logger.info("\n步骤 3: 测试结果统计")
        logger.info("=" * 60)
        
        success_count = sum(1 for r in results if r['success'])
        failed_count = len(results) - success_count
        
        logger.info(f"总任务数: {len(results)}")
        logger.info(f"成功数: {success_count}")
        logger.info(f"失败数: {failed_count}")
        logger.info(f"成功率: {success_count/len(results)*100:.1f}%")
        
        # 第四步：清理环境（可选）
        if Config.AUTO_CLOSE_BROWSER:
            logger.info("\n步骤 4: 清理测试环境")
            logger.info("-" * 60)
            
            # 删除环境
            if self.created_profiles:
                self.delete_profiles(self.created_profiles)
            
            # 删除代理
            if self.created_proxies:
                self.delete_proxies(self.created_proxies)
        
        logger.info("\n" + "=" * 60)
        logger.info("并发测试完成！")
        logger.info("=" * 60)


if __name__ == "__main__":
    logger.info("并发测试工具")
    logger.info(f"当前环境: {Config.ENV}")
    logger.info(f"AdsPower API: {Config.ADSPOWER_API_URL}")
    logger.info("")
    
    # ==================== 配置区 ====================
    TEST_URL = "https://health.baidu.com/m/detail/ar_2366617956693492811"
    NUM_TASKS = 3  # 并发任务数
    MAX_WORKERS = 3  # 最大并发执行数（建议不超过3）
    # =====================================================
    
    tester = ConcurrentTester(
        test_url=TEST_URL,
        max_workers=MAX_WORKERS
    )
    
    tester.run_concurrent_test(num_tasks=NUM_TASKS)