#!/usr/bin/env python3
"""
优化的网站截图脚本 - 使用 Playwright
- 更快的加载速度
- 截图后自动删除文件
- 不保存临时文件
"""

import asyncio
from playwright.async_api import async_playwright
import sys
import os
import tempfile
from datetime import datetime

async def take_screenshot(url, output_path=None, auto_delete=True):
    """
    截取网站截图
    """
    temp_file = None
    
    try:
        if output_path is None:
            # 使用临时文件
            fd, temp_file = tempfile.mkstemp(suffix='.png', prefix='screenshot_')
            os.close(fd)
            output_path = temp_file
        else:
            # 如果指定了输出路径，确保目录存在
            os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        
        print(f"正在打开: {url}")
        
        async with async_playwright() as p:
            # 启动浏览器（优化参数）
            browser = await p.chromium.launch(
                headless=True,
                args=[
                    "--no-sandbox",
                    "--disable-setuid-sandbox",
                    "--disable-dev-shm-usage",
                    "--disable-gpu",
                    "--disable-software-rasterizer",
                    "--disable-extensions",
                    "--no-first-run",
                    "--no-default-browser-check",
                    "--disable-background-networking",
                    "--disable-default-apps",
                    "--disable-sync",
                    "--disable-translate",
                    "--metrics-recording-only",
                    "--disable-prompt-on-repost",
                    "--disable-hang-monitor",
                    "--disable-client-side-phishing-detection",
                    "--disable-component-update",
                    "--disable-domain-reliability",
                ]
            )
            
            # 创建页面
            page = await browser.new_page(
                viewport={"width": 1920, "height": 1080}
            )
            
            # 访问网站（等待页面加载完成）
            await page.goto(url, wait_until="load", timeout=30000)
            # 额外等待一下让 JS 执行
            await asyncio.sleep(2)
            
            # 截图（不全页，更快）
            await page.screenshot(path=output_path, full_page=False)
            print(f"截图已保存: {output_path}")
            
            # 关闭浏览器
            await browser.close()
        
        return output_path
        
    finally:
        # 如果是临时文件且设置了自动删除，则删除
        if auto_delete and temp_file and os.path.exists(temp_file):
            try:
                os.unlink(temp_file)
                print(f"临时文件已删除: {temp_file}")
            except:
                pass

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("使用方法: python screenshot_website.py <url> [output_path]")
        print("示例: python screenshot_website.py https://kirkify.net")
        sys.exit(1)
    
    url = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else None
    
    try:
        asyncio.run(take_screenshot(url, output_path))
    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()