"""
视觉测试模块
截图对比、视觉回归测试
"""
import hashlib
from typing import Optional, List, Tuple
from pathlib import Path
from playwright.async_api import Page
from loguru import logger
from dataclasses import dataclass
from datetime import datetime
import base64


@dataclass
class ScreenshotDiff:
    """截图差异"""
    baseline_path: str
    current_path: str
    diff_path: Optional[str] = None
    diff_percentage: float = 0.0
    is_match: bool = True
    pixel_diff: int = 0


class ScreenshotManager:
    """
    截图管理器
    管理截图存储和命名
    """
    
    def __init__(self, base_dir: str = "./screenshots"):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self.baseline_dir = self.base_dir / "baseline"
        self.current_dir = self.base_dir / "current"
        self.diff_dir = self.base_dir / "diff"
        
        for d in [self.baseline_dir, self.current_dir, self.diff_dir]:
            d.mkdir(parents=True, exist_ok=True)
    
    def get_baseline_path(self, name: str) -> Path:
        """获取基线路径"""
        return self.baseline_dir / f"{name}.png"
    
    def get_current_path(self, name: str) -> Path:
        """获取当前截图路径"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return self.current_dir / f"{name}_{timestamp}.png"
    
    def get_diff_path(self, name: str) -> Path:
        """获取差异图路径"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return self.diff_dir / f"{name}_diff_{timestamp}.png"
    
    def list_baselines(self) -> List[str]:
        """列出所有基线截图"""
        return [f.stem for f in self.baseline_dir.glob("*.png")]
    
    def has_baseline(self, name: str) -> bool:
        """检查是否存在基线"""
        return self.get_baseline_path(name).exists()


class VisualTester:
    """
    视觉测试器
    支持截图对比和视觉回归测试
    """
    
    def __init__(self, page: Page, base_dir: str = "./screenshots"):
        self.page = page
        self.manager = ScreenshotManager(base_dir)
        self.threshold: float = 0.01  # 默认差异阈值1%
    
    async def capture(
        self,
        name: str,
        selector: Optional[str] = None,
        full_page: bool = False,
    ) -> str:
        """
        截图
        
        Args:
            name: 截图名称
            selector: 元素选择器（可选）
            full_page: 是否全页面
            
        Returns:
            str: 截图路径
        """
        path = str(self.manager.get_current_path(name))
        
        if selector:
            await self.page.locator(selector).screenshot(path=path)
        else:
            await self.page.screenshot(path=path, full_page=full_page)
        
        logger.info(f"截图已保存: {path}")
        return path
    
    async def set_baseline(
        self,
        name: str,
        selector: Optional[str] = None,
        full_page: bool = False,
    ) -> str:
        """
        设置基线截图
        
        Args:
            name: 截图名称
            selector: 元素选择器
            full_page: 是否全页面
            
        Returns:
            str: 基线路径
        """
        path = str(self.manager.get_baseline_path(name))
        
        if selector:
            await self.page.locator(selector).screenshot(path=path)
        else:
            await self.page.screenshot(path=path, full_page=full_page)
        
        logger.info(f"基线截图已保存: {path}")
        return path
    
    async def compare(
        self,
        name: str,
        selector: Optional[str] = None,
        full_page: bool = False,
        threshold: Optional[float] = None,
    ) -> ScreenshotDiff:
        """
        对比截图与基线
        
        Args:
            name: 截图名称
            selector: 元素选择器
            full_page: 是否全页面
            threshold: 差异阈值
            
        Returns:
            ScreenshotDiff: 对比结果
        """
        baseline_path = self.manager.get_baseline_path(name)
        
        # 如果没有基线，自动创建
        if not baseline_path.exists():
            logger.warning(f"基线不存在，自动创建: {name}")
            await self.set_baseline(name, selector, full_page)
            return ScreenshotDiff(
                baseline_path=str(baseline_path),
                current_path=str(baseline_path),
                is_match=True,
            )
        
        # 获取当前截图
        current_path = self.manager.get_current_path(name)
        
        if selector:
            await self.page.locator(selector).screenshot(path=str(current_path))
        else:
            await self.page.screenshot(path=str(current_path), full_page=full_page)
        
        # 对比截图
        diff = await self._compare_images(
            str(baseline_path),
            str(current_path),
            threshold or self.threshold,
        )
        
        return diff
    
    async def _compare_images(
        self,
        baseline_path: str,
        current_path: str,
        threshold: float,
    ) -> ScreenshotDiff:
        """
        对比两张图片
        
        Args:
            baseline_path: 基线路径
            current_path: 当前截图路径
            threshold: 差异阈值
            
        Returns:
            ScreenshotDiff: 对比结果
        """
        try:
            from PIL import Image
            import numpy as np
            
            # 读取图片
            baseline = Image.open(baseline_path)
            current = Image.open(current_path)
            
            # 确保尺寸一致
            if baseline.size != current.size:
                logger.warning(f"截图尺寸不一致: {baseline.size} vs {current.size}")
                current = current.resize(baseline.size)
            
            # 转换为数组
            baseline_arr = np.array(baseline)
            current_arr = np.array(current)
            
            # 计算差异
            diff_arr = np.abs(baseline_arr.astype(int) - current_arr.astype(int))
            pixel_diff = np.sum(diff_arr > 0)
            total_pixels = baseline_arr.shape[0] * baseline_arr.shape[1]
            
            # 计算差异百分比
            diff_percentage = pixel_diff / total_pixels
            is_match = diff_percentage <= threshold
            
            # 生成差异图
            diff_path = None
            if not is_match:
                diff_path = str(self.manager.get_diff_path(Path(current_path).stem))
                diff_img = Image.fromarray(np.uint8(diff_arr * 255 / diff_arr.max()))
                diff_img.save(diff_path)
            
            result = ScreenshotDiff(
                baseline_path=baseline_path,
                current_path=current_path,
                diff_path=diff_path,
                diff_percentage=diff_percentage * 100,
                is_match=is_match,
                pixel_diff=int(pixel_diff),
            )
            
            if is_match:
                logger.info(f"截图对比通过: 差异={diff_percentage*100:.2f}%")
            else:
                logger.warning(
                    f"截图对比失败: 差异={diff_percentage*100:.2f}%, "
                    f"阈值={threshold*100:.2f}%"
                )
            
            return result
            
        except ImportError:
            logger.warning("PIL未安装，使用简化对比")
            return await self._simple_compare(baseline_path, current_path, threshold)
    
    async def _simple_compare(
        self,
        baseline_path: str,
        current_path: str,
        threshold: float,
    ) -> ScreenshotDiff:
        """简化对比（不需要PIL）"""
        import hashlib
        
        def get_file_hash(path: str) -> str:
            with open(path, 'rb') as f:
                return hashlib.md5(f.read()).hexdigest()
        
        baseline_hash = get_file_hash(baseline_path)
        current_hash = get_file_hash(current_path)
        
        is_match = baseline_hash == current_hash
        
        return ScreenshotDiff(
            baseline_path=baseline_path,
            current_path=current_path,
            diff_percentage=0.0 if is_match else 100.0,
            is_match=is_match,
        )
    
    async def update_baseline(self, name: str) -> bool:
        """
        更新基线截图
        
        Args:
            name: 截图名称
            
        Returns:
            bool: 是否成功
        """
        current_files = list(self.manager.current_dir.glob(f"{name}_*.png"))
        
        if not current_files:
            logger.error(f"未找到当前截图: {name}")
            return False
        
        # 使用最新的当前截图
        latest = sorted(current_files)[-1]
        baseline_path = self.manager.get_baseline_path(name)
        
        import shutil
        shutil.copy(latest, baseline_path)
        
        logger.info(f"基线已更新: {name}")
        return True
    
    async def clear_current(self):
        """清除当前截图"""
        import shutil
        for f in self.manager.current_dir.glob("*.png"):
            f.unlink()
        logger.info("已清除当前截图")
    
    async def clear_diff(self):
        """清除差异图"""
        for f in self.manager.diff_dir.glob("*.png"):
            f.unlink()
        logger.info("已清除差异图")


class ScreenshotAnnotator:
    """
    截图标注器
    在截图上添加标注
    """
    
    def __init__(self, page: Page):
        self.page = page
    
    async def highlight_element(
        self,
        selector: str,
        color: str = "red",
        label: Optional[str] = None,
    ) -> bytes:
        """
        标注元素截图
        
        Args:
            selector: 选择器
            color: 标注颜色
            label: 标签文字
            
        Returns:
            bytes: 截图数据
        """
        # 注入标注脚本
        await self.page.evaluate(f"""
            () => {{
                const element = document.querySelector('{selector}');
                if (element) {{
                    const rect = element.getBoundingClientRect();
                    const overlay = document.createElement('div');
                    overlay.style.cssText = `
                        position: fixed;
                        top: {rect.top}px;
                        left: {rect.left}px;
                        width: {rect.width}px;
                        height: {rect.height}px;
                        border: 3px solid {color};
                        background: transparent;
                        pointer-events: none;
                        z-index: 999999;
                    `;
                    overlay.id = 'screenshot-annotation';
                    document.body.appendChild(overlay);
                    
                    {'// 添加标签' if label else ''}
                    {f'''
                    const labelDiv = document.createElement('div');
                    labelDiv.style.cssText = `
                        position: fixed;
                        top: {rect.top - 25}px;
                        left: {rect.left}px;
                        background: {color};
                        color: white;
                        padding: 2px 8px;
                        font-size: 12px;
                        font-family: sans-serif;
                        border-radius: 3px;
                        z-index: 999999;
                    `;
                    labelDiv.textContent = '{label}';
                    document.body.appendChild(labelDiv);
                    ''' if label else ''}
                }}
            }}
        """)
        
        # 截图
        screenshot = await self.page.screenshot()
        
        # 移除标注
        await self.page.evaluate("""
            () => {
                const overlays = document.querySelectorAll('[id^="screenshot-annotation"], [style*="border: 3px solid"]');
                overlays.forEach(el => el.remove());
            }
        """)
        
        return screenshot
    
    async def add_watermark(
        self,
        text: str = "Screenshot",
        position: str = "bottom-right",
        opacity: float = 0.5,
    ) -> bytes:
        """
        添加水印
        
        Args:
            text: 水印文字
            position: 位置
            opacity: 透明度
            
        Returns:
            bytes: 截图数据
        """
        positions = {
            "top-left": "top: 10px; left: 10px;",
            "top-right": "top: 10px; right: 10px;",
            "bottom-left": "bottom: 10px; left: 10px;",
            "bottom-right": "bottom: 10px; right: 10px;",
            "center": "top: 50%; left: 50%; transform: translate(-50%, -50%);",
        }
        
        pos_style = positions.get(position, positions["bottom-right"])
        
        await self.page.evaluate(f"""
            () => {{
                const watermark = document.createElement('div');
                watermark.style.cssText = `
                    position: fixed;
                    {pos_style}
                    background: rgba(0, 0, 0, {opacity});
                    color: white;
                    padding: 5px 15px;
                    font-size: 14px;
                    font-family: sans-serif;
                    border-radius: 3px;
                    z-index: 999999;
                    pointer-events: none;
                `;
                watermark.textContent = '{text}';
                watermark.id = 'screenshot-watermark';
                document.body.appendChild(watermark);
            }}
        """)
        
        screenshot = await self.page.screenshot()
        
        # 移除水印
        await self.page.evaluate("""
            () => {
                const watermark = document.getElementById('screenshot-watermark');
                if (watermark) watermark.remove();
            }
        """)
        
        return screenshot
    
    async def add_timestamp(self, format: str = "%Y-%m-%d %H:%M:%S") -> bytes:
        """
        添加时间戳
        
        Args:
            format: 时间格式
            
        Returns:
            bytes: 截图数据
        """
        timestamp = datetime.now().strftime(format)
        return await self.add_watermark(timestamp, "bottom-left")


class VisualReporter:
    """
    视觉测试报告生成器
    """
    
    def __init__(self, output_dir: str = "./reports"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.results: List[ScreenshotDiff] = []
    
    def add_result(self, result: ScreenshotDiff):
        """添加对比结果"""
        self.results.append(result)
    
    def generate_html_report(self, name: str = "visual_report") -> str:
        """
        生成HTML报告
        
        Args:
            name: 报告名称
            
        Returns:
            str: 报告路径
        """
        html = """
<!DOCTYPE html>
<html>
<head>
    <title>Visual Test Report</title>
    <style>
        body { font-family: Arial, sans-serif; padding: 20px; }
        .summary { margin-bottom: 20px; padding: 15px; background: #f5f5f5; border-radius: 5px; }
        .result { margin-bottom: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 5px; }
        .pass { border-left: 5px solid #4CAF50; }
        .fail { border-left: 5px solid #f44336; }
        .images { display: flex; gap: 20px; margin-top: 10px; }
        .images img { max-width: 300px; border: 1px solid #ddd; }
    </style>
</head>
<body>
    <h1>Visual Test Report</h1>
    <div class="summary">
        <p>Total: {total}</p>
        <p>Passed: {passed}</p>
        <p>Failed: {failed}</p>
    </div>
    <div class="results">
        {results_html}
    </div>
</body>
</html>
        """
        
        passed = sum(1 for r in self.results if r.is_match)
        failed = len(self.results) - passed
        
        results_html = ""
        for result in self.results:
            status = "pass" if result.is_match else "fail"
            results_html += f"""
        <div class="result {status}">
            <h3>{Path(result.current_path).stem}</h3>
            <p>Status: {'PASS' if result.is_match else 'FAIL'}</p>
            <p>Diff: {result.diff_percentage:.2f}%</p>
            <div class="images">
                <div>
                    <p>Baseline:</p>
                    <img src="{result.baseline_path}" />
                </div>
                <div>
                    <p>Current:</p>
                    <img src="{result.current_path}" />
                </div>
                {'<div><p>Diff:</p><img src="' + result.diff_path + '" /></div>' if result.diff_path else ''}
            </div>
        </div>
            """
        
        report_path = self.output_dir / f"{name}.html"
        report_path.write_text(html.format(
            total=len(self.results),
            passed=passed,
            failed=failed,
            results_html=results_html,
        ))
        
        logger.info(f"报告已生成: {report_path}")
        return str(report_path)
