ai_wht_wechat/backend/error_screenshot.py

"""
错误截图保存工具
当发生错误时自动截图并保存，便于问题排查
"""
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional
from playwright.async_api import Page


# 截图保存目录
SCREENSHOT_DIR = Path("error_screenshots")
SCREENSHOT_DIR.mkdir(exist_ok=True)


async def save_error_screenshot(
    page: Optional[Page],
    error_type: str,
    error_message: str = "",
    prefix: str = ""
) -> Optional[str]:
    """
    保存错误截图
    
    Args:
        page: Playwright 页面对象
        error_type: 错误类型（如：login_failed, send_code_failed, publish_failed等）
        error_message: 错误信息（可选，会添加到日志）
        prefix: 文件名前缀（可选）
    
    Returns:
        截图文件路径，失败返回None
    """
    if not page:
        print("[错误截图] 页面对象为空，无法截图", file=sys.stderr)
        return None
    
    try:
        # 检查页面状态
        try:
            current_url = page.url
            print(f"[错误截图] 当前URL: {current_url}", file=sys.stderr)
            
            # 检查是否是空白页
            if current_url in ['about:blank', '', 'data:,']:
                print(f"[错误截图] 警告: 当前页面为空白页，截图可能没有内容", file=sys.stderr)
            
            # 等待页面稳定
            await page.wait_for_load_state('domcontentloaded', timeout=3000)
        except Exception as state_error:
            print(f"[错误截图] 检查页面状态失败: {str(state_error)}", file=sys.stderr)
        
        # 生成文件名：年月日时分秒_错误类型.png
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 清理错误类型字符串（移除特殊字符）
        safe_error_type = "".join(c for c in error_type if c.isalnum() or c in ('_', '-'))
        
        # 组合文件名
        if prefix:
            filename = f"{prefix}_{timestamp}_{safe_error_type}.png"
        else:
            filename = f"{timestamp}_{safe_error_type}.png"
        
        filepath = SCREENSHOT_DIR / filename
        
        # 截图（添加超时和全页截图）
        await page.screenshot(
            path=str(filepath), 
            full_page=True,
            timeout=10000  # 10秒超时
        )
        
        # 检查截图文件大小
        file_size = filepath.stat().st_size
        print(f"[错误截图] 已保存: {filepath} (大小: {file_size} bytes)", file=sys.stderr)
        
        # 如果文件太小（小于5KB），可能是空白截图
        if file_size < 5120:
            print(f"[错误截图] 警告: 截图文件过小 ({file_size} bytes)，可能为空白页面", file=sys.stderr)
        
        if error_message:
            print(f"[错误截图] 错误信息: {error_message}", file=sys.stderr)
        
        # 返回文件路径
        return str(filepath)
    
    except Exception as e:
        print(f"[错误截图] 截图失败: {str(e)}", file=sys.stderr)
        return None


def cleanup_old_screenshots(days: int = 7):
    """
    清理旧的错误截图
    
    Args:
        days: 保留最近几天的截图，默认7天
    """
    try:
        import time
        current_time = time.time()
        cutoff_time = current_time - (days * 24 * 60 * 60)
        
        deleted_count = 0
        for file in SCREENSHOT_DIR.glob("*.png"):
            if file.stat().st_mtime < cutoff_time:
                file.unlink()
                deleted_count += 1
        
        if deleted_count > 0:
            print(f"[错误截图] 已清理 {deleted_count} 个超过 {days} 天的旧截图", file=sys.stderr)
    
    except Exception as e:
        print(f"[错误截图] 清理旧截图失败: {str(e)}", file=sys.stderr)


async def save_screenshot_with_html(
    page: Optional[Page],
    error_type: str,
    error_message: str = "",
    prefix: str = ""
) -> tuple[Optional[str], Optional[str]]:
    """
    保存错误截图和HTML源码（用于深度调试）
    
    Args:
        page: Playwright 页面对象
        error_type: 错误类型
        error_message: 错误信息（可选）
        prefix: 文件名前缀（可选）
    
    Returns:
        (截图路径, HTML路径)，失败返回(None, None)
    """
    if not page:
        return None, None
    
    try:
        # 检查页面状态
        try:
            current_url = page.url
            print(f"[错误截图] 当前URL: {current_url}", file=sys.stderr)
            
            if current_url in ['about:blank', '', 'data:,']:
                print(f"[错误截图] 警告: 当前页面为空白页", file=sys.stderr)
            
            # 等待页面稳定
            await page.wait_for_load_state('domcontentloaded', timeout=3000)
        except Exception as state_error:
            print(f"[错误截图] 检查页面状态失败: {str(state_error)}", file=sys.stderr)
        
        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_error_type = "".join(c for c in error_type if c.isalnum() or c in ('_', '-'))
        
        if prefix:
            base_filename = f"{prefix}_{timestamp}_{safe_error_type}"
        else:
            base_filename = f"{timestamp}_{safe_error_type}"
        
        # 保存截图
        screenshot_path = SCREENSHOT_DIR / f"{base_filename}.png"
        await page.screenshot(
            path=str(screenshot_path), 
            full_page=True,
            timeout=10000
        )
        
        # 检查截图文件大小
        screenshot_size = screenshot_path.stat().st_size
        if screenshot_size < 5120:
            print(f"[错误截图] 警告: 截图文件过小 ({screenshot_size} bytes)", file=sys.stderr)
        
        # 保存HTML
        html_path = SCREENSHOT_DIR / f"{base_filename}.html"
        html_content = await page.content()
        with open(html_path, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        html_size = html_path.stat().st_size
        print(f"[错误截图] 已保存截图: {screenshot_path} ({screenshot_size} bytes)", file=sys.stderr)
        print(f"[错误截图] 已保存HTML: {html_path} ({html_size} bytes)", file=sys.stderr)
        if error_message:
            print(f"[错误截图] 错误信息: {error_message}", file=sys.stderr)
        
        return str(screenshot_path), str(html_path)
    
    except Exception as e:
        print(f"[错误截图] 保存截图和HTML失败: {str(e)}", file=sys.stderr)
        return None, None