#!/home/hz1/miniconda3/envs/openclaw/bin/python3.12
"""
Web Capture API - 网页截图与代码提取服务
使用 agent-browser (Rust + Playwright) 或 Playwright
"""

import os
import json
import asyncio
import subprocess
import tempfile
import uuid
import time
from pathlib import Path
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
from datetime import datetime

# 检查 Playwright 是否可用
try:
    from playwright.async_api import async_playwright
    PLAYWRIGHT_AVAILABLE = True
except ImportError:
    PLAYWRIGHT_AVAILABLE = False

# 检查 agent-browser 是否可用
def check_agent_browser():
    try:
        result = subprocess.run(
            ["which", "agent-browser"],
            capture_output=True,
            text=True
        )
        return result.returncode == 0
    except:
        return False

AGENT_BROWSER_AVAILABLE = check_agent_browser()

app = Flask(__name__)
CORS(app)

# 配置
CAPTURE_DIR = Path(tempfile.gettempdir()) / "web_captures"
CAPTURE_DIR.mkdir(exist_ok=True)


class AgentBrowserSession:
    """agent-browser 会话管理器"""
    
    def __init__(self, session_id=None):
        self.session_id = session_id or str(uuid.uuid4())[:8]
        self.base_cmd = ["agent-browser", "--session", self.session_id]
    
    def run(self, cmd, timeout=60000):
        """执行 agent-browser 命令"""
        full_cmd = self.base_cmd + cmd
        try:
            result = subprocess.run(
                full_cmd,
                capture_output=True,
                text=True,
                timeout=timeout / 1000
            )
            return result.returncode == 0, result.stdout, result.stderr
        except subprocess.TimeoutExpired:
            return False, "", "Command timeout"
        except Exception as e:
            return False, "", str(e)
    
    def open(self, url):
        """打开网页"""
        return self.run(["open", url], timeout=30000)
    
    def set_viewport(self, width, height):
        """设置视口大小"""
        return self.run(["set", "viewport", str(width), str(height)])
    
    def screenshot(self, output_path, full_page=False):
        """截图"""
        cmd = ["screenshot", str(output_path)]
        if full_page:
            cmd.append("--full")
        return self.run(cmd, timeout=30000)
    
    def get_html(self):
        """获取HTML"""
        success, output, error = self.run(
            ["eval", "document.documentElement.outerHTML"],
            timeout=10000
        )
        return success, output, error
    
    def scroll_down(self, pixels=800):
        """向下滚动"""
        return self.run(["scroll", "down", str(pixels)])
    
    def wait(self, ms):
        """等待"""
        return self.run(["wait", str(ms)])
    
    def close(self):
        """关闭浏览器"""
        try:
            self.run(["close"])
        except:
            pass


def capture_with_agent_browser(
    url: str,
    action: str = "screenshot",
    scroll_times: int = 0,
    scroll_delay: int = 1000,
    full_page: bool = False,
    viewport: dict = None,
    wait_time: int = 2000
):
    """
    使用 agent-browser 捕获网页
    """
    session = AgentBrowserSession()
    temp_file = None
    
    try:
        # 设置视口
        if viewport:
            session.set_viewport(viewport.get("width", 1920), viewport.get("height", 1080))
        
        # 打开网页
        success, stdout, stderr = session.open(url)
        if not success:
            return {"success": False, "error": f"Failed to open URL: {stderr}"}
        
        # 等待页面加载
        session.wait(wait_time)
        
        # 滚动加载动态内容
        if scroll_times > 0:
            for i in range(scroll_times):
                session.scroll_down(800)
                session.wait(scroll_delay)
        
        # 执行操作
        if action == "screenshot":
            temp_file = CAPTURE_DIR / f"{session.session_id}.png"
            success, _, error = session.screenshot(temp_file, full_page=full_page)
            
            if success and temp_file.exists():
                return {"success": True, "file_path": str(temp_file)}
            else:
                return {"success": False, "error": f"Screenshot failed: {error}"}
        
        elif action == "html":
            success, html, error = session.get_html()
            if success:
                return {"success": True, "html": html}
            else:
                return {"success": False, "error": f"Get HTML failed: {error}"}
        
        else:
            return {"success": False, "error": f"Unknown action: {action}"}
    
    except Exception as e:
        return {"success": False, "error": str(e)}
    
    finally:
        session.close()


async def capture_with_playwright(
    url: str,
    action: str = "screenshot",
    scroll_times: int = 0,
    scroll_delay: int = 1000,
    full_page: bool = False,
    viewport: dict = None,
    wait_time: int = 2000
):
    """
    使用 Playwright 捕获网页
    """
    if not PLAYWRIGHT_AVAILABLE:
        return {"success": False, "error": "Playwright not installed"}
    
    session_id = str(uuid.uuid4())[:8]
    
    try:
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True)
            
            viewport_settings = viewport or {"width": 1920, "height": 1080}
            
            context = await browser.new_context(
                viewport={
                    "width": viewport_settings.get("width", 1920),
                    "height": viewport_settings.get("height", 1080)
                },
                user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
            )
            
            page = await context.new_page()
            
            try:
                await page.goto(url, wait_until="networkidle", timeout=30000)
            except:
                await page.goto(url, wait_until="domcontentloaded", timeout=30000)
            
            if wait_time > 0:
                await page.wait_for_timeout(wait_time)
            
            if scroll_times > 0:
                for i in range(scroll_times):
                    await page.evaluate("window.scrollBy(0, 800)")
                    await page.wait_for_timeout(scroll_delay)
                await page.wait_for_timeout(500)
            
            result = {}
            
            if action == "screenshot":
                temp_file = CAPTURE_DIR / f"{session_id}.png"
                await page.screenshot(path=str(temp_file), full_page=full_page)
                result = {"success": True, "file_path": str(temp_file)}
            
            elif action == "html":
                html = await page.content()
                result = {"success": True, "html": html}
            
            else:
                result = {"success": False, "error": f"Unknown action: {action}"}
            
            await browser.close()
            return result
    
    except Exception as e:
        return {"success": False, "error": str(e)}


def capture_webpage(
    url: str,
    action: str = "screenshot",
    scroll_times: int = 0,
    scroll_delay: int = 1000,
    full_page: bool = False,
    viewport: dict = None,
    wait_time: int = 2000,
    backend: str = "auto"
):
    """
    捕获网页
    
    Args:
        backend: "auto", "agent-browser", "playwright"
    """
    # 选择后端
    if backend == "auto":
        if AGENT_BROWSER_AVAILABLE:
            backend = "agent-browser"
        elif PLAYWRIGHT_AVAILABLE:
            backend = "playwright"
        else:
            return {"success": False, "error": "No browser backend available. Install agent-browser or playwright."}
    
    # 使用对应后端
    if backend == "agent-browser":
        if not AGENT_BROWSER_AVAILABLE:
            return {"success": False, "error": "agent-browser not available"}
        return capture_with_agent_browser(
            url, action, scroll_times, scroll_delay, full_page, viewport, wait_time
        )
    
    elif backend == "playwright":
        if not PLAYWRIGHT_AVAILABLE:
            return {"success": False, "error": "Playwright not available"}
        
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            result = loop.run_until_complete(
                capture_with_playwright(
                    url, action, scroll_times, scroll_delay, full_page, viewport, wait_time
                )
            )
            loop.close()
            return result
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    else:
        return {"success": False, "error": f"Unknown backend: {backend}"}


@app.route('/')
def index():
    """主页"""
    return render_template('index.html')


@app.route('/api')
def api_info():
    """API信息"""
    return jsonify({
        "service": "Web Capture API",
        "version": "2.0.0",
        "backends": {
            "agent-browser": "available" if AGENT_BROWSER_AVAILABLE else "unavailable",
            "playwright": "available" if PLAYWRIGHT_AVAILABLE else "unavailable"
        },
        "endpoints": {
            "/api/capture": "POST - Capture webpage screenshot or HTML",
            "/health": "GET - Health check"
        }
    })


@app.route('/health')
def health():
    """健康检查"""
    status = "healthy" if (AGENT_BROWSER_AVAILABLE or PLAYWRIGHT_AVAILABLE) else "degraded"
    return jsonify({
        "status": status,
        "agent-browser": "installed" if AGENT_BROWSER_AVAILABLE else "not found",
        "playwright": "installed" if PLAYWRIGHT_AVAILABLE else "not found",
        "timestamp": datetime.now().isoformat()
    })


@app.route('/api/capture', methods=['POST'])
def capture():
    """
    捕获网页接口
    
    请求体:
    {
        "url": "https://example.com",
        "action": "screenshot" | "html",
        "scroll_times": 0,
        "scroll_delay": 1000,
        "full_page": false,
        "viewport": {"width": 1920, "height": 1080},
        "wait_time": 2000,
        "backend": "auto" | "agent-browser" | "playwright"
    }
    """
    data = request.get_json()
    
    if not data:
        return jsonify({"success": False, "error": "No JSON data provided"}), 400
    
    url = data.get("url")
    if not url:
        return jsonify({"success": False, "error": "URL is required"}), 400
    
    # 添加协议前缀
    if not url.startswith(("http://", "https://")):
        url = "https://" + url
    
    action = data.get("action", "screenshot")
    scroll_times = int(data.get("scroll_times", 0))
    scroll_delay = int(data.get("scroll_delay", 1000))
    full_page = bool(data.get("full_page", False))
    viewport = data.get("viewport")
    wait_time = int(data.get("wait_time", 2000))
    backend = data.get("backend", "auto")
    
    result = capture_webpage(
        url=url,
        action=action,
        scroll_times=scroll_times,
        scroll_delay=scroll_delay,
        full_page=full_page,
        viewport=viewport,
        wait_time=wait_time,
        backend=backend
    )
    
    if not result["success"]:
        return jsonify(result), 400
    
    if action == "screenshot":
        file_path = result["file_path"]
        return send_file(file_path, mimetype='image/png')
    
    elif action == "html":
        return jsonify(result)


if __name__ == '__main__':
    print("🌐 Web Capture API Starting...")
    print(f"📁 Capture directory: {CAPTURE_DIR}")
    print(f"🔧 Backends:")
    print(f"   - agent-browser: {'✅' if AGENT_BROWSER_AVAILABLE else '❌'}")
    print(f"   - playwright: {'✅' if PLAYWRIGHT_AVAILABLE else '❌'}")
    
    if not AGENT_BROWSER_AVAILABLE and not PLAYWRIGHT_AVAILABLE:
        print("\n⚠️  没有可用的浏览器后端，请安装其中一个:")
        print("   agent-browser: npm install -g agent-browser && agent-browser install")
        print("   playwright: pip install playwright && playwright install chromium")
    
    print("\n🚀 Server running on http://0.0.0.0:5000")
    app.run(host='0.0.0.0', port=5000, debug=True)