"""
截图辅助工具模块

提供测试截图功能，支持多种截图模式。
"""

import os
from pathlib import Path
from typing import Optional
from datetime import datetime
from playwright.sync_api import Page


class ScreenshotHelper:
    """截图辅助工具"""

    def __init__(self, screenshot_dir: str = "reports/screenshots"):
        self.screenshot_dir = Path(screenshot_dir)
        self.screenshot_dir.mkdir(parents=True, exist_ok=True)

    def take_screenshot(
        self,
        page: Page,
        name: str,
        full_page: bool = False,
        selector: Optional[str] = None,
    ) -> str:
        """
        截取页面截图

        Args:
            page: Playwright页面对象
            name: 截图文件名
            full_page: 是否截取整个页面
            selector: 特定元素选择器

        Returns:
            截图文件路径
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{name}_{timestamp}.png"
        filepath = self.screenshot_dir / filename

        try:
            if selector:
                # 截取特定元素
                element = page.locator(selector)
                element.screenshot(path=str(filepath))
            else:
                # 截取页面
                page.screenshot(path=str(filepath), full_page=full_page)

            print(f"📸 截图已保存: {filepath}")
            return str(filepath)
        except Exception as e:
            print(f"❌ 截图失败: {e}")
            return ""

    def take_screenshot_on_failure(
        self, page: Page, test_name: str, full_page: bool = True
    ) -> str:
        """
        测试失败时截取截图

        Args:
            page: Playwright页面对象
            test_name: 测试名称
            full_page: 是否截取整个页面

        Returns:
            截图文件路径
        """
        return self.take_screenshot(
            page, f"{test_name}_failed", full_page=full_page
        )

    def take_comparison_screenshot(
        self, page: Page, name: str, baseline_dir: str = "screenshots/baseline"
    ) -> tuple:
        """
        截取对比截图

        Args:
            page: Playwright页面对象
            name: 截图名称
            baseline_dir: 基线截图目录

        Returns:
            (当前截图路径, 基线截图路径)
        """
        current_path = self.take_screenshot(page, name, full_page=True)

        baseline_path = Path(baseline_dir) / f"{name}.png"

        return current_path, str(baseline_path)

    def cleanup_old_screenshots(self, days: int = 7) -> int:
        """
        清理旧截图

        Args:
            days: 保留天数

        Returns:
            删除的文件数量
        """
        import time

        deleted_count = 0
        cutoff_time = time.time() - (days * 24 * 60 * 60)

        for file_path in self.screenshot_dir.glob("*.png"):
            if file_path.stat().st_mtime < cutoff_time:
                file_path.unlink()
                deleted_count += 1

        print(f"🗑️ 已清理 {deleted_count} 个旧截图文件")
        return deleted_count