import gradio as gr
from playwright.sync_api import sync_playwright
from PIL import Image
import io
import subprocess
import sys
import os

# ==================== 自动安装浏览器 ====================
def install_browsers():
    print("Downloading Playwright browsers...")
    try:
        subprocess.run(["playwright", "install", "chromium", "--with-deps"], check=True)
        print("Browser installation completed!")
    except Exception as e:
        print(f"Browser install failed: {e}")

if "SPACE_ID" in os.environ:   # HF Space 环境
    install_browsers()

# ==================== 主函数 ====================
def run_automation(url: str, action: str = "screenshot"):
    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            page = browser.new_page(viewport={"width": 1280, "height": 720})
            page.goto(url, wait_until="networkidle", timeout=60000)
            
            if action == "screenshot":
                screenshot_bytes = page.screenshot(full_page=True)
                browser.close()
                image = Image.open(io.BytesIO(screenshot_bytes))
                return "✅ 截图成功！", image
            else:
                title = page.title()
                browser.close()
                return f"📄 页面标题：{title}", None
                
    except Exception as e:
        return f"❌ 错误：{str(e)}", None

# ==================== Gradio 界面 ====================
with gr.Blocks(title="网页截图工具") as demo:
    gr.Markdown("# 🌐 网页自动化演示（Playwright）")
    
    with gr.Row():
        url_input = gr.Textbox(
            label="网址", 
            placeholder="https://huggingface.co",
            value="https://example.com",
            scale=4
        )
        action = gr.Radio(["screenshot", "title"], label="操作", value="screenshot")
    
    btn = gr.Button("🚀 执行", variant="primary")
    
    with gr.Row():
        output_text = gr.Textbox(label="状态")
        output_image = gr.Image(label="截图结果")

    btn.click(
        fn=run_automation,
        inputs=[url_input, action],
        outputs=[output_text, output_image]
    )

demo.launch()