# -*- coding: utf-8 -*-
"""
ChromeDriver 自动管理模块
- 检测Chrome浏览器安装
- 检测ChromeDriver
- 自动下载匹配版本的ChromeDriver
"""

import os
import sys
import platform
import subprocess
import zipfile
import shutil
import json
import re
from urllib.request import urlopen, urlretrieve
from urllib.error import URLError

# 下载进度回调类型
# callback(current, total, status_text)


class ChromeDriverManager:
    """ChromeDriver管理器"""
    
    # ChromeDriver下载镜像站点
    MIRROR_URLS = {
        'official': 'https://storage.googleapis.com/chrome-for-testing-public',
        'taobao': 'https://registry.npmmirror.com/-/binary/chromedriver',
        'huawei': 'https://mirrors.huawei.com/chrome-driver'
    }
    
    def __init__(self):
        self.system = platform.system().lower()
        self.machine = platform.machine().lower()
        
        # 确定平台标识
        if self.system == 'windows':
            self.platform = 'win' + ('64' if self.machine in ['amd64', 'x86_64'] else '32')
        elif self.system == 'darwin':
            self.platform = 'mac-' + ('arm64' if self.machine == 'arm64' else 'x64')
        else:
            self.platform = 'linux64'
        
        # ChromeDriver保存目录
        self.driver_dir = os.path.join(os.path.dirname(__file__), 'drivers')
        if not os.path.exists(self.driver_dir):
            os.makedirs(self.driver_dir)
    
    def get_chrome_version(self):
        """获取Chrome浏览器版本"""
        version = None
        
        try:
            if self.system == 'windows':
                # 方法1：注册表
                try:
                    import winreg
                    key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, 
                        r"Software\Google\Chrome\BLBeacon")
                    version, _ = winreg.QueryValueEx(key, "version")
                    winreg.CloseKey(key)
                except:
                    pass
                
                # 方法2：默认安装路径
                if not version:
                    paths = [
                        os.path.join(os.environ.get('PROGRAMFILES', ''), 
                                    'Google', 'Chrome', 'Application', 'chrome.exe'),
                        os.path.join(os.environ.get('PROGRAMFILES(X86)', ''), 
                                    'Google', 'Chrome', 'Application', 'chrome.exe'),
                        os.path.join(os.environ.get('LOCALAPPDATA', ''), 
                                    'Google', 'Chrome', 'Application', 'chrome.exe'),
                    ]
                    for path in paths:
                        if os.path.exists(path):
                            try:
                                result = subprocess.run(
                                    [path, '--version'],
                                    capture_output=True, text=True, timeout=5
                                )
                                match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                                if match:
                                    version = match.group(1)
                                    break
                            except:
                                pass
            
            elif self.system == 'darwin':
                # macOS
                paths = [
                    '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome',
                    '/Applications/Chromium.app/Contents/MacOS/Chromium'
                ]
                for path in paths:
                    if os.path.exists(path):
                        try:
                            result = subprocess.run(
                                [path, '--version'],
                                capture_output=True, text=True, timeout=5
                            )
                            match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                            if match:
                                version = match.group(1)
                                break
                        except:
                            pass
            
            else:
                # Linux
                result = subprocess.run(
                    ['google-chrome', '--version'],
                    capture_output=True, text=True, timeout=5
                )
                match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                if match:
                    version = match.group(1)
        
        except Exception as e:
            print(f"获取Chrome版本失败: {e}")
        
        return version
    
    def is_chrome_installed(self):
        """检查Chrome浏览器是否已安装"""
        return self.get_chrome_version() is not None
    
    def get_chromedriver_path(self):
        """获取ChromeDriver路径"""
        if self.system == 'windows':
            driver_name = 'chromedriver.exe'
        else:
            driver_name = 'chromedriver'
        
        # 检查drivers目录
        driver_path = os.path.join(self.driver_dir, driver_name)
        if os.path.exists(driver_path):
            return driver_path
        
        # 检查系统PATH
        try:
            result = subprocess.run(
                ['where' if self.system == 'windows' else 'which', 'chromedriver'],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0:
                return result.stdout.strip().split('\n')[0]
        except:
            pass
        
        return None
    
    def get_chromedriver_version(self):
        """获取ChromeDriver版本"""
        driver_path = self.get_chromedriver_path()
        if not driver_path:
            return None
        
        try:
            result = subprocess.run(
                [driver_path, '--version'],
                capture_output=True, text=True, timeout=5
            )
            match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
            if match:
                return match.group(1)
        except:
            pass
        
        return None
    
    def is_chromedriver_compatible(self):
        """检查ChromeDriver版本是否与Chrome兼容"""
        chrome_version = self.get_chrome_version()
        driver_version = self.get_chromedriver_version()
        
        if not chrome_version or not driver_version:
            return False
        
        # 主版本号必须匹配
        chrome_major = chrome_version.split('.')[0]
        driver_major = driver_version.split('.')[0]
        
        return chrome_major == driver_major
    
    def get_download_url(self, version, callback=None):
        """获取ChromeDriver下载URL"""
        # 提取主版本号
        major_version = version.split('.')[0]
        
        # 尝试官方源
        urls_to_try = [
            # Chrome 115+ 使用新格式
            f"{self.MIRROR_URLS['official']}/{version}/{self.platform}/chromedriver-{self.platform}.zip",
            # 旧格式
            f"{self.MIRROR_URLS['taobao']}/{version}/chromedriver_{self.platform}.zip",
        ]
        
        for url in urls_to_try:
            try:
                if callback:
                    callback(0, 0, f"尝试下载源: {url[:50]}...")
                
                # 测试URL是否可访问
                response = urlopen(url, timeout=10)
                if response.getcode() == 200:
                    return url
            except:
                continue
        
        return None
    
    def download_chromedriver(self, version=None, callback=None):
        """下载ChromeDriver
        
        Args:
            version: Chrome版本，如果不提供则自动检测
            callback: 进度回调函数 callback(current, total, status)
        
        Returns:
            (success, message, driver_path)
        """
        try:
            # 获取Chrome版本
            if not version:
                version = self.get_chrome_version()
            
            if not version:
                return False, "未检测到Chrome浏览器，请先安装Chrome", None
            
            if callback:
                callback(0, 0, f"检测到Chrome版本: {version}")
            
            # 获取下载URL
            download_url = self.get_download_url(version, callback)
            
            if not download_url:
                return False, f"未找到Chrome {version}对应的ChromeDriver下载地址", None
            
            if callback:
                callback(0, 0, f"开始下载ChromeDriver...")
            
            # 下载文件
            zip_path = os.path.join(self.driver_dir, 'chromedriver.zip')
            
            def download_progress(block_num, block_size, total_size):
                if callback and total_size > 0:
                    current = block_num * block_size
                    percent = min(100, int(current * 100 / total_size))
                    callback(current, total_size, f"下载中... {percent}%")
            
            try:
                urlretrieve(download_url, zip_path, download_progress)
            except Exception as e:
                return False, f"下载失败: {str(e)}", None
            
            if callback:
                callback(0, 0, "下载完成，正在解压...")
            
            # 解压文件
            try:
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall(self.driver_dir)
                
                # 查找解压后的chromedriver
                driver_name = 'chromedriver.exe' if self.system == 'windows' else 'chromedriver'
                driver_path = None
                
                for root, dirs, files in os.walk(self.driver_dir):
                    for file in files:
                        if file == driver_name:
                            driver_path = os.path.join(root, file)
                            break
                
                if not driver_path:
                    # 尝试查找包含chromedriver的文件
                    for root, dirs, files in os.walk(self.driver_dir):
                        for file in files:
                            if 'chromedriver' in file.lower():
                                # 移动到drivers目录
                                src = os.path.join(root, file)
                                driver_path = os.path.join(self.driver_dir, driver_name)
                                shutil.move(src, driver_path)
                                break
                
                # 清理临时文件
                os.remove(zip_path)
                
                # 设置执行权限（Linux/Mac）
                if self.system != 'windows' and driver_path:
                    os.chmod(driver_path, 0o755)
                
                if callback:
                    callback(1, 1, f"ChromeDriver安装成功: {driver_path}")
                
                return True, f"ChromeDriver {version} 安装成功", driver_path
            
            except Exception as e:
                return False, f"解压失败: {str(e)}", None
        
        except Exception as e:
            return False, f"安装过程出错: {str(e)}", None
    
    def ensure_chromedriver(self, auto_download=True, callback=None):
        """确保ChromeDriver可用
        
        Args:
            auto_download: 是否自动下载
            callback: 进度回调
        
        Returns:
            (success, message, driver_path)
        """
        # 检查Chrome是否安装
        if not self.is_chrome_installed():
            return False, "未检测到Chrome浏览器，请先安装Chrome浏览器", None
        
        chrome_version = self.get_chrome_version()
        
        # 检查ChromeDriver是否存在且兼容
        if self.is_chromedriver_compatible():
            driver_path = self.get_chromedriver_path()
            return True, f"ChromeDriver已就绪 (Chrome {chrome_version})", driver_path
        
        # ChromeDriver不存在或不兼容
        if auto_download:
            if callback:
                callback(0, 0, f"ChromeDriver版本不匹配，正在自动下载...")
            return self.download_chromedriver(chrome_version, callback)
        else:
            driver_version = self.get_chromedriver_version()
            if driver_version:
                return False, f"ChromeDriver版本不匹配 (Chrome: {chrome_version}, Driver: {driver_version})", None
            else:
                return False, f"未找到ChromeDriver (Chrome版本: {chrome_version})", None


# 全局实例
_chrome_driver_manager = None

def get_chrome_driver_manager():
    """获取ChromeDriver管理器单例"""
    global _chrome_driver_manager
    if _chrome_driver_manager is None:
        _chrome_driver_manager = ChromeDriverManager()
    return _chrome_driver_manager


def check_and_install_chromedriver(parent_widget=None):
    """检查并安装ChromeDriver（带UI提示）
    
    Args:
        parent_widget: 父窗口，用于显示对话框
    
    Returns:
        (success, driver_path)
    """
    from PyQt6.QtWidgets import (
        QDialog, QVBoxLayout, QLabel, QProgressBar, QPushButton, 
        QMessageBox, QApplication
    )
    from PyQt6.QtCore import Qt, QThread, pyqtSignal
    
    manager = get_chrome_driver_manager()
    
    # 先快速检查
    result = manager.ensure_chromedriver(auto_download=False)
    
    if result[0]:
        return True, result[2]
    
    # 需要下载
    if parent_widget is None:
        # 没有父窗口，尝试无界面下载
        success, msg, path = manager.download_chromedriver()
        return success, path
    
    # 显示下载对话框
    class DownloadThread(QThread):
        progress = pyqtSignal(int, int, str)
        finished = pyqtSignal(bool, str, str)
        
        def run(self):
            success, msg, path = manager.download_chromedriver(
                callback=lambda c, t, s: self.progress.emit(c, t, s)
            )
            self.finished.emit(success, msg, path or '')
    
    dialog = QDialog(parent_widget)
    dialog.setWindowTitle("安装ChromeDriver")
    dialog.setFixedSize(400, 180)
    dialog.setWindowFlags(dialog.windowFlags() & ~Qt.WindowType.WindowCloseButtonHint)
    
    layout = QVBoxLayout(dialog)
    layout.setSpacing(15)
    
    # 状态标签
    status_label = QLabel("正在检测ChromeDriver...")
    status_label.setWordWrap(True)
    layout.addWidget(status_label)
    
    # 进度条
    progress_bar = QProgressBar()
    progress_bar.setRange(0, 100)
    progress_bar.setValue(0)
    layout.addWidget(progress_bar)
    
    # 详细信息
    detail_label = QLabel("")
    detail_label.setStyleSheet("color: #666; font-size: 11px;")
    layout.addWidget(detail_label)
    
    # 取消按钮（初始隐藏）
    cancel_btn = QPushButton("取消")
    cancel_btn.setVisible(False)
    layout.addWidget(cancel_btn)
    
    download_thread = DownloadThread()
    
    def on_progress(current, total, status):
        status_label.setText(status)
        detail_label.setText(f"已下载: {current // 1024} KB / {total // 1024} KB" if total > 0 else status)
        if total > 0:
            progress_bar.setValue(int(current * 100 / total))
    
    def on_finished(success, msg, path):
        progress_bar.setValue(100 if success else 0)
        if success:
            status_label.setText("ChromeDriver安装成功!")
            status_label.setStyleSheet("color: #27ae60; font-weight: bold;")
            dialog.accept()
        else:
            status_label.setText(f"安装失败: {msg}")
            status_label.setStyleSheet("color: #e74c3c;")
            cancel_btn.setText("关闭")
            cancel_btn.setVisible(True)
    
    download_thread.progress.connect(on_progress)
    download_thread.finished.connect(on_finished)
    
    # 启动下载
    download_thread.start()
    
    result = dialog.exec()
    
    if result == QDialog.DialogCode.Accepted:
        return True, manager.get_chromedriver_path()
    else:
        return False, None


if __name__ == "__main__":
    # 测试
    manager = ChromeDriverManager()
    print(f"Chrome版本: {manager.get_chrome_version()}")
    print(f"ChromeDriver路径: {manager.get_chromedriver_path()}")
    print(f"ChromeDriver版本: {manager.get_chromedriver_version()}")
    print(f"兼容性: {manager.is_chromedriver_compatible()}")
