#!/usr/bin/env python3
"""
视觉比对器模块
用于比较截图与设计图的视觉差异，支持PNG、JPG和SVG格式
"""

import cv2
import numpy as np
from pathlib import Path
from skimage.metrics import structural_similarity as ssim
from PIL import Image
from config import Config

# 条件导入cairosvg，如果失败则使用备用方案
CAIROSVG_AVAILABLE = False
try:
    import cairosvg
    CAIROSVG_AVAILABLE = True
    print("✅ CairoSVG库可用，支持完整SVG处理")
except (ImportError, OSError) as e:
    print(f"⚠️  CairoSVG库不可用: {e}")
    print("💡 将使用备用的SVG处理方案")

# 条件导入selenium，用于SVG渲染
SELENIUM_AVAILABLE = False
try:
    from selenium import webdriver
    from selenium.webdriver.chrome.options import Options
    from selenium.webdriver.chrome.service import Service
    from selenium.webdriver.common.by import By
    from webdriver_manager.chrome import ChromeDriverManager
    import base64
    SELENIUM_AVAILABLE = True
    print("✅ Selenium库可用，支持浏览器SVG渲染")
except ImportError as e:
    print(f"⚠️  Selenium库不可用: {e}")
    print("💡 将跳过浏览器SVG渲染方案")


class VisualComparator:
    """视觉比对器"""
    
    def __init__(self):
        self.comparison_results = []
        self.config = Config.VISUAL_COMPARISON
    
    def compare_images(self, screenshot_path, design_path, output_path=None):
        """比较截图和设计图"""
        try:
            # 读取图像
            screenshot = self._load_image(screenshot_path)
            design = self._load_image(design_path)
            
            if screenshot is None or design is None:
                print(f"❌ 无法加载图像: {screenshot_path} 或 {design_path}")
                return None
            
            # 智能调整图像尺寸，保持宽高比
            screenshot_processed, design_processed = self._smart_resize_images(screenshot, design)
            
            # 转换为灰度图
            gray1 = cv2.cvtColor(screenshot_processed, cv2.COLOR_BGR2GRAY)
            gray2 = cv2.cvtColor(design_processed, cv2.COLOR_BGR2GRAY)
            
            # 计算结构相似性
            similarity_score, diff_image = ssim(gray1, gray2, full=True)
            diff_image = (diff_image * 255).astype(np.uint8)
            
            # 查找差异区域
            diff_regions = self._find_diff_regions(diff_image)
            
            # 生成比对结果图
            if output_path:
                self._generate_comparison_image(
                    screenshot_processed, design_processed, diff_image, 
                    diff_regions, similarity_score, output_path
                )
            
            result = {
                'similarity_score': similarity_score,
                'diff_regions': diff_regions,
                'total_diff_area': sum(r['area'] for r in diff_regions),
                'diff_percentage': (1 - similarity_score) * 100,
                'meets_threshold': similarity_score >= self.config['similarity_threshold'],
                'original_screenshot_size': screenshot.shape[:2],
                'original_design_size': design.shape[:2],
                'processed_size': screenshot_processed.shape[:2]
            }
            
            self.comparison_results.append(result)
            return result
            
        except Exception as e:
            print(f"❌ 图像比对失败: {e}")
            return None
    
    def _load_image(self, image_path):
        """加载图像，支持PNG、JPG和SVG格式"""
        image_path = Path(image_path)
        
        if not image_path.exists():
            print(f"❌ 文件不存在: {image_path}")
            return None
        
        # 检查文件格式
        if image_path.suffix.lower() == '.svg':
            return self._load_svg_image(image_path)
        else:
            # 加载常规图像格式
            img = cv2.imread(str(image_path))
            if img is None:
                print(f"❌ 无法读取图像文件: {image_path}")
            return img
    
    def _load_svg_image(self, svg_path):
        """加载SVG图像并转换为OpenCV格式"""
        import io
        
        if CAIROSVG_AVAILABLE:
            # 优先使用cairosvg
            try:
                # 将SVG转换为PNG字节数据
                png_data = cairosvg.svg2png(url=str(svg_path))
                
                # 使用PIL加载PNG数据
                pil_image = Image.open(io.BytesIO(png_data))
                
                # 转换为OpenCV格式
                opencv_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
                
                return opencv_image
                
            except Exception as e:
                print(f"❌ CairoSVG处理失败: {e}")
        
        # 备用方案1: 尝试使用wand (ImageMagick)
        try:
            from wand.image import Image as WandImage
            
            with WandImage(filename=str(svg_path)) as img:
                img.format = 'png'
                blob = img.make_blob()
                
            pil_image = Image.open(io.BytesIO(blob))
            opencv_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
            
            print("✅ 使用Wand处理SVG成功")
            return opencv_image
            
        except ImportError:
            print("⚠️  Wand库不可用")
        except Exception as e:
            print(f"❌ Wand处理失败: {e}")
        
        # 备用方案2: 尝试使用reportlab和svglib
        try:
            from reportlab.graphics import renderPM
            from svglib.svglib import renderSVG
            
            drawing = renderSVG(str(svg_path))
            pil_image = renderPM.drawToPIL(drawing)
            opencv_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
            
            print("✅ 使用ReportLab处理SVG成功")
            return opencv_image
            
        except ImportError:
            print("⚠️  ReportLab/svglib库不可用")
        except Exception as e:
            print(f"❌ ReportLab处理失败: {e}")
        
        # 备用方案3: 尝试使用Selenium浏览器渲染
        if SELENIUM_AVAILABLE:
            try:
                result = self._render_svg_with_selenium(svg_path)
                if result is not None:
                    print("✅ 使用Selenium浏览器渲染SVG成功")
                    return result
            except Exception as e:
                print(f"❌ Selenium渲染失败: {e}")
        
        # 备用方案4: 简单的文本提示图像
        print(f"⚠️  无法处理SVG文件: {svg_path}")
        print("💡 建议安装以下任一库来支持SVG:")
        print("   - pip install cairosvg (推荐，但需要Cairo库)")
        print("   - pip install Wand (需要ImageMagick)")
        print("   - pip install reportlab svglib")
        
        # 创建一个提示图像
        placeholder_img = np.ones((400, 600, 3), dtype=np.uint8) * 240
        cv2.putText(placeholder_img, "SVG Not Supported", (150, 180), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (100, 100, 100), 2)
        cv2.putText(placeholder_img, "Install SVG library", (150, 220), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (100, 100, 100), 2)
        
        return placeholder_img
    
    def _render_svg_with_selenium(self, svg_path):
        """使用Selenium浏览器渲染SVG"""
        try:
            # 配置Chrome选项
            chrome_options = Options()
            chrome_options.add_argument('--headless')  # 无头模式
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--disable-web-security')
            chrome_options.add_argument('--allow-running-insecure-content')
            chrome_options.add_argument('--disable-extensions')
            chrome_options.add_argument('--window-size=1200,800')
            
            # 创建WebDriver
            service = Service(ChromeDriverManager().install())
            driver = webdriver.Chrome(service=service, options=chrome_options)
            
            try:
                # 读取SVG文件内容
                with open(svg_path, 'r', encoding='utf-8') as f:
                    svg_content = f.read()
                
                # 创建HTML页面包含SVG
                html_content = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <style>
        body {{ margin: 0; padding: 20px; background: white; }}
        svg {{ max-width: 100%; height: auto; display: block; }}
    </style>
</head>
<body>
    {svg_content}
</body>
</html>"""
                
                # 创建临时HTML文件
                import tempfile
                import os
                
                with tempfile.NamedTemporaryFile(mode='w', suffix='.html', delete=False, encoding='utf-8') as f:
                    f.write(html_content)
                    temp_html_path = f.name
                
                try:
                    # 使用file:// URL加载本地文件
                    file_url = f"file:///{temp_html_path.replace(os.sep, '/')}"
                    driver.get(file_url)
                    
                    # 等待页面加载
                    import time
                    time.sleep(3)
                    
                    # 获取SVG元素
                    svg_element = driver.find_element(By.TAG_NAME, "svg")
                    
                    # 截图SVG元素
                    screenshot = svg_element.screenshot_as_png
                    
                    # 转换为OpenCV格式
                    import io
                    pil_image = Image.open(io.BytesIO(screenshot))
                    opencv_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
                    
                    return opencv_image
                    
                finally:
                    # 清理临时文件
                    try:
                        os.unlink(temp_html_path)
                    except:
                        pass
                
            finally:
                driver.quit()
                
        except Exception as e:
            print(f"Selenium渲染SVG时出错: {e}")
            return None
    
    def _smart_resize_images(self, img1, img2):
        """智能调整图像尺寸，保持宽高比"""
        h1, w1 = img1.shape[:2]
        h2, w2 = img2.shape[:2]
        
        # 计算宽高比
        ratio1 = w1 / h1
        ratio2 = w2 / h2
        
        # 如果宽高比相近（差异小于10%），则按较小的尺寸调整
        if abs(ratio1 - ratio2) / max(ratio1, ratio2) < 0.1:
            # 选择较小的尺寸作为目标尺寸
            target_w = min(w1, w2)
            target_h = min(h1, h2)
            
            img1_resized = cv2.resize(img1, (target_w, target_h))
            img2_resized = cv2.resize(img2, (target_w, target_h))
            
            return img1_resized, img2_resized
        
        # 如果宽高比差异较大，则保持原始比例，按最大公共区域对齐
        else:
            # 计算最大公共尺寸，保持各自的宽高比
            if ratio1 > ratio2:  # img1更宽
                # 以img2的高度为基准
                new_h = min(h1, h2)
                new_w1 = int(new_h * ratio1)
                new_w2 = int(new_h * ratio2)
            else:  # img2更宽或相等
                # 以img1的高度为基准
                new_h = min(h1, h2)
                new_w1 = int(new_h * ratio1)
                new_w2 = int(new_h * ratio2)
            
            # 调整图像尺寸
            img1_resized = cv2.resize(img1, (new_w1, new_h))
            img2_resized = cv2.resize(img2, (new_w2, new_h))
            
            # 如果宽度不同，则创建相同宽度的画布，居中放置
            max_w = max(new_w1, new_w2)
            
            if new_w1 != max_w:
                canvas1 = np.ones((new_h, max_w, 3), dtype=np.uint8) * 255
                offset = (max_w - new_w1) // 2
                canvas1[:, offset:offset+new_w1] = img1_resized
                img1_resized = canvas1
            
            if new_w2 != max_w:
                canvas2 = np.ones((new_h, max_w, 3), dtype=np.uint8) * 255
                offset = (max_w - new_w2) // 2
                canvas2[:, offset:offset+new_w2] = img2_resized
                img2_resized = canvas2
            
            return img1_resized, img2_resized
    
    def _resize_to_match(self, img1, img2):
        """调整图像尺寸匹配（保留原方法作为备用）"""
        h2, w2 = img2.shape[:2]
        return cv2.resize(img1, (w2, h2))
    
    def _find_diff_regions(self, diff_image):
        """查找差异区域"""
        threshold = self.config['diff_threshold']
        min_area = self.config['min_diff_area']
        
        # 二值化差异图像
        _, binary = cv2.threshold(diff_image, threshold, 255, cv2.THRESH_BINARY)
        
        # 查找轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        regions = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if area > min_area:  # 过滤小差异
                x, y, w, h = cv2.boundingRect(contour)
                regions.append({
                    'x': x, 'y': y, 'width': w, 'height': h,
                    'area': area,
                    'center_x': x + w//2, 'center_y': y + h//2
                })
        
        return regions
    
    def _generate_comparison_image(self, img1, img2, diff_img, regions, score, output_path):
        """生成比对结果图像"""
        # 创建拼接图像
        h, w = img1.shape[:2]
        result_img = np.zeros((h, w*3, 3), dtype=np.uint8)
        
        # 放置原图像
        result_img[:, :w] = img1
        result_img[:, w:2*w] = img2
        result_img[:, 2*w:] = cv2.cvtColor(diff_img, cv2.COLOR_GRAY2BGR)
        
        # 标记差异区域
        for region in regions:
            x, y, rw, rh = region['x'], region['y'], region['width'], region['height']
            # 在截图上标记
            cv2.rectangle(result_img, (x, y), (x+rw, y+rh), (0, 0, 255), 2)
            # 在设计图上标记
            cv2.rectangle(result_img, (w+x, y), (w+x+rw, y+rh), (0, 0, 255), 2)
        
        # 添加文本信息
        font = cv2.FONT_HERSHEY_SIMPLEX
        
        # 相似度信息
        similarity_text = f"Similarity: {score:.2%}"
        cv2.putText(result_img, similarity_text, (10, 30), font, 1, (255, 255, 255), 2)
        
        # 差异统计
        diff_count = len(regions)
        diff_text = f"Differences: {diff_count}"
        cv2.putText(result_img, diff_text, (10, 70), font, 0.8, (255, 255, 255), 2)
        
        # 标签
        cv2.putText(result_img, "Screenshot", (10, h-10), font, 0.7, (255, 255, 255), 2)
        cv2.putText(result_img, "Design", (w+10, h-10), font, 0.7, (255, 255, 255), 2)
        cv2.putText(result_img, "Difference", (2*w+10, h-10), font, 0.7, (255, 255, 255), 2)
        
        # 保存结果
        cv2.imwrite(str(output_path), result_img)
        print(f"📊 比对结果已保存: {output_path}")
    
    def get_comparison_summary(self):
        """获取比对结果摘要"""
        if not self.comparison_results:
            return None
        
        latest_result = self.comparison_results[-1]
        
        summary = {
            'total_comparisons': len(self.comparison_results),
            'latest_similarity': latest_result['similarity_score'],
            'latest_diff_count': len(latest_result['diff_regions']),
            'meets_threshold': latest_result['meets_threshold'],
            'average_similarity': sum(r['similarity_score'] for r in self.comparison_results) / len(self.comparison_results)
        }
        
        return summary
    
    def find_design_files(self, design_dir):
        """查找设计目录中的所有支持格式文件"""
        design_dir = Path(design_dir)
        supported_formats = self.config['supported_formats']
        
        design_files = []
        for format_ext in supported_formats:
            design_files.extend(design_dir.glob(f"*{format_ext}"))
        
        # 优先返回SVG文件
        svg_files = [f for f in design_files if f.suffix.lower() == '.svg']
        other_files = [f for f in design_files if f.suffix.lower() != '.svg']
        
        return svg_files + other_files
    
    def batch_compare(self, screenshot_path, design_dir, output_dir=None):
        """批量比对截图与设计目录中的所有文件"""
        design_files = self.find_design_files(design_dir)
        
        if not design_files:
            print(f"❌ 在 {design_dir} 中未找到支持的设计文件")
            return []
        
        results = []
        for design_file in design_files:
            print(f"🔍 比对设计文件: {design_file.name}")
            
            output_path = None
            if output_dir:
                output_dir = Path(output_dir)
                output_dir.mkdir(exist_ok=True)
                output_path = output_dir / f"comparison_{design_file.stem}.png"
            
            result = self.compare_images(screenshot_path, design_file, output_path)
            if result:
                result['design_file'] = str(design_file)
                results.append(result)
        
        return results