from typing import Callable, Any
from functools import wraps


def retry_on_failure(max_retries: int = 3, delay: int = 1000):
    """失败重试装饰器

    Args:
        max_retries: 最大重试次数
        delay: 重试延迟（毫秒）
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            import time

            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        time.sleep(delay / 1000)
                    else:
                        raise last_exception

        return wrapper

    return decorator


def handle_test_failure(test_name: str):
    """测试失败处理装饰器

    Args:
        test_name: 测试名称
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            try:
                return func(*args, **kwargs)
            except Exception as e:
                import os
                from datetime import datetime

                os.makedirs("screenshots", exist_ok=True)
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"screenshots/failure_{test_name}_{timestamp}.png"

                if "page" in kwargs:
                    kwargs["page"].screenshot(path=filename)

                print(f"测试失败: {test_name}")
                print(f"错误信息: {str(e)}")
                print(f"截图已保存: {filename}")

                raise e

        return wrapper

    return decorator