from playwright.sync_api import Page, Locator
from typing import Optional
from pathlib import Path
import os


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

    def __init__(self, page: Page, screenshot_dir: str = "screenshots"):
        """初始化截图辅助工具

        Args:
            page: Playwright页面对象
            screenshot_dir: 截图保存目录
        """
        self.page = page
        self.screenshot_dir = screenshot_dir
        self._ensure_screenshot_dir()

    def _ensure_screenshot_dir(self) -> None:
        """确保截图目录存在"""
        Path(self.screenshot_dir).mkdir(parents=True, exist_ok=True)

    def take_full_page_screenshot(self, name: str, timeout: int = 30000) -> str:
        """截取整个页面

        Args:
            name: 截图文件名（不含扩展名）
            timeout: 页面加载超时时间（毫秒）

        Returns:
            截图文件路径
        """
        self.page.wait_for_load_state("networkidle", timeout=timeout)
        file_path = os.path.join(self.screenshot_dir, f"{name}.png")
        self.page.screenshot(path=file_path, full_page=True)
        return file_path

    def take_viewport_screenshot(self, name: str, timeout: int = 30000) -> str:
        """截取当前视口

        Args:
            name: 截图文件名（不含扩展名）
            timeout: 页面加载超时时间（毫秒）

        Returns:
            截图文件路径
        """
        self.page.wait_for_load_state("networkidle", timeout=timeout)
        file_path = os.path.join(self.screenshot_dir, f"{name}.png")
        self.page.screenshot(path=file_path, full_page=False)
        return file_path

    def take_element_screenshot(self, locator: Locator, name: str, timeout: int = 10000) -> str:
        """截取指定元素

        Args:
            locator: 元素定位器
            name: 截图文件名（不含扩展名）
            timeout: 元素等待超时时间（毫秒）

        Returns:
            截图文件路径
        """
        element = locator.wait_for(timeout=timeout)
        file_path = os.path.join(self.screenshot_dir, f"{name}.png")
        element.screenshot(path=file_path)
        return file_path

    def take_element_screenshot_by_selector(
        self, selector: str, name: str, timeout: int = 10000
    ) -> str:
        """通过选择器截取元素

        Args:
            selector: CSS选择器
            name: 截图文件名（不含扩展名）
            timeout: 元素等待超时时间（毫秒）

        Returns:
            截图文件路径
        """
        locator = self.page.locator(selector)
        return self.take_element_screenshot(locator, name, timeout)

    def take_screenshot_with_mask(
        self, name: str, mask_selectors: list, timeout: int = 30000
    ) -> str:
        """截取页面并遮蔽指定元素

        Args:
            name: 截图文件名（不含扩展名）
            mask_selectors: 需要遮蔽的元素选择器列表
            timeout: 页面加载超时时间（毫秒）

        Returns:
            截图文件路径
        """
        self.page.wait_for_load_state("networkidle", timeout=timeout)
        file_path = os.path.join(self.screenshot_dir, f"{name}.png")

        masks = []
        for selector in mask_selectors:
            element = self.page.locator(selector)
            if element.is_visible():
                masks.append(element)

        self.page.screenshot(path=file_path, full_page=True, mask=masks)
        return file_path

    def take_screenshot_on_failure(self, test_name: str, step_name: Optional[str] = None) -> str:
        """测试失败时截图

        Args:
            test_name: 测试名称
            step_name: 步骤名称（可选）

        Returns:
            截图文件路径
        """
        if step_name:
            file_name = f"{test_name}_{step_name}_failed"
        else:
            file_name = f"{test_name}_failed"

        return self.take_full_page_screenshot(file_name)

    def take_screenshot_series(self, base_name: str, count: int, delay: int = 1000) -> list:
        """连续截取多张截图

        Args:
            base_name: 截图基础名称
            count: 截图数量
            delay: 每次截图之间的延迟（毫秒）

        Returns:
            截图文件路径列表
        """
        file_paths = []
        for i in range(count):
            file_name = f"{base_name}_{i + 1}"
            file_path = self.take_viewport_screenshot(file_name)
            file_paths.append(file_path)
            if i < count - 1:
                self.page.wait_for_timeout(delay)
        return file_paths

    def take_screenshot_before_and_after(self, action, name: str) -> tuple:
        """在操作前后截图

        Args:
            action: 要执行的操作函数
            name: 截图基础名称

        Returns:
            (操作前截图路径, 操作后截图路径)
        """
        before_path = self.take_viewport_screenshot(f"{name}_before")
        action()
        after_path = self.take_viewport_screenshot(f"{name}_after")
        return (before_path, after_path)

    def capture_visual_diff(self, name: str, expected_path: str) -> tuple:
        """捕获视觉差异

        Args:
            name: 截图文件名
            expected_path: 期望截图路径

        Returns:
            (当前截图路径, 差异截图路径)
        """
        current_path = self.take_full_page_screenshot(name)

        if not os.path.exists(expected_path):
            return (current_path, None)

        return (current_path, None)