#!/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
    # 尝试导入 stealth 反检测
    try:
        from playwright_stealth import stealth_async
        STEALTH_AVAILABLE = True
    except ImportError:
        STEALTH_AVAILABLE = False
        print("⚠️  playwright-stealth 未安装，反爬虫能力降低")
except ImportError:
    PLAYWRIGHT_AVAILABLE = False
    STEALTH_AVAILABLE = False
    print("⚠️  Playwright 未安装，请运行: pip install playwright && playwright install chromium")

# 检查 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)

# agent-browser socket 目录
os.environ.setdefault("AGENT_BROWSER_SOCKET_DIR", "/tmp/agent-browser-sockets")
Path(os.environ["AGENT_BROWSER_SOCKET_DIR"]).mkdir(exist_ok=True, parents=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
        # 设置必要的环境变量
        env = os.environ.copy()
        env['AGENT_BROWSER_SOCKET_DIR'] = '/tmp/agent-browser-sockets'
        env['AGENT_BROWSER_USER_AGENT'] = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        try:
            result = subprocess.run(
                full_cmd,
                capture_output=True,
                text=True,
                timeout=timeout / 1000,
                env=env
            )
            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):
        """打开网页"""
        # 添加反爬虫检测的 User-Agent
        env = os.environ.copy()
        env['AGENT_BROWSER_USER_AGENT'] = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        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:
            # 检查是否是反爬虫拦截
            if "403" in stdout or "Access Denied" in stdout:
                return {"success": False, "error": "网站反爬虫拦截 (403)，建议使用 Playwright 后端或手动添加请求头"}
            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 捕获网页
    Playwright 有更强的反爬虫能力
    """
    if not PLAYWRIGHT_AVAILABLE:
        return {"success": False, "error": "Playwright not installed. Run: pip install playwright && playwright install chromium"}
    
    session_id = str(uuid.uuid4())[:8]
    
    try:
        async with async_playwright() as p:
            # 启动浏览器，添加反检测配置
            # 尝试使用真实 Chrome（如果安装了）
            try:
                # 首先尝试使用真实 Chrome 浏览器
                browser = await p.chromium.launch(
                    headless=True,
                    channel='chrome',  # 使用真实 Chrome 而不是 Chromium
                    args=[
                        '--disable-blink-features=AutomationControlled',
                        '--disable-dev-shm-usage',
                        '--no-sandbox',
                        '--ignore-certificate-errors',
                    ]
                )
            except:
                # 如果没有 Chrome，使用 Chromium
                browser = await p.chromium.launch(
                    headless=True,
                    args=[
                        '--disable-blink-features=AutomationControlled',
                        '--disable-dev-shm-usage',
                        '--no-sandbox',
                        '--ignore-certificate-errors',
                    ]
                )
            
            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 (KHTML, like Gecko) Chrome/120.0.6099.109 Safari/537.36',
                # 添加更多浏览器特征
                locale='en-US',
                timezone_id='America/New_York',
                color_scheme='light',
                has_touch=False,
                is_mobile=False,
                java_script_enabled=True,
                accept_downloads=True,
            )
            
            # 注入反检测脚本
            await context.add_init_script("""
                // 移除 webdriver 标记
                Object.defineProperty(navigator, 'webdriver', {
                    get: () => undefined
                });
                
                // 添加插件
                Object.defineProperty(navigator, 'plugins', {
                    get: () => {
                        return [
                            {name: 'Chrome PDF Plugin', filename: 'internal-pdf-viewer', description: 'Portable Document Format'},
                            {name: 'Chrome PDF Viewer', filename: 'mhjfbmdgncjokjjpdkfnmlglggk', description: ''},
                            {name: 'Native Client', filename: 'internal-nacl-plugin', description: ''}
                        ];
                    }
                });
                
                // 语言列表
                Object.defineProperty(navigator, 'languages', {
                    get: () => ['en-US', 'en', 'zh-CN']
                });
                
                // 平台
                Object.defineProperty(navigator, 'platform', {
                    get: () => 'Win32'
                });
                
                // 硬件并发
                Object.defineProperty(navigator, 'hardwareConcurrency', {
                    get: () => 8
                });
                
                // 设备内存
                Object.defineProperty(navigator, 'deviceMemory', {
                    get: () => 8
                });
                
                // Chrome 对象
                window.chrome = {
                    runtime: {
                        PlatformOs: {MAC: 'mac', WIN: 'win', ANDROID: 'android', CROS: 'cros', LINUX: 'linux', OPENBSD: 'openbsd'},
                        PlatformArch: {ARM: 'arm', X86_32: 'x86-32', X86_64: 'x86-64'},
                        PlatformNaclArch: {ARM: 'arm', X86_32: 'x86-32', X86_64: 'x86-64'},
                        Request: function() {},
                        connect: function() {},
                        sendMessage: function() {}
                    },
                    loadTimes: function() {},
                    csi: function() {},
                    app: {}
                };
            """)
            
            page = await context.new_page()
            
            # 应用 stealth 反检测（如果可用）
            if STEALTH_AVAILABLE:
                await stealth_async(page)
            
            try:
                await page.goto(url, wait_until="domcontentloaded", timeout=30000)
            except Exception as e:
                # 如果 domcontentloaded 超时，尝试 commit
                try:
                    await page.goto(url, wait_until="commit", timeout=30000)
                except:
                    pass
            
            # 固定等待时间（用于验证、加载等过程）
            if wait_time > 0:
                await page.wait_for_timeout(wait_time)
            
            # 等待页面基本稳定
            try:
                await page.wait_for_load_state("networkidle", timeout=5000)
            except:
                # networkidle 超时没关系，继续执行
                pass
            
            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)}


async def capture_with_cdp(
    url_hint: str = "",
    action: str = "screenshot",
    cdp_port: int = 9222
):
    """
    使用 CDP 连接到已打开的 Chrome 浏览器
    用于方案三：手动验证后自动截图
    """
    if not PLAYWRIGHT_AVAILABLE:
        return {"success": False, "error": "Playwright not installed"}
    
    try:
        async with async_playwright() as p:
            # 连接到已运行的 Chrome
            browser = await p.chromium.connect_over_cdp(f"http://localhost:{cdp_port}")
            
            # 获取所有页面
            contexts = browser.contexts
            result_data = []
            
            for context in contexts:
                for page in context.pages:
                    current_url = page.url
                    title = await page.title()
                    
                    # 如果 URL 包含提示词或者是唯一页面
                    if url_hint in current_url or not url_hint:
                        if action == "screenshot":
                            temp_file = CAPTURE_DIR / f"cdp_capture_{uuid.uuid4()[:8]}.png"
                            await page.screenshot(path=str(temp_file), full_page=True)
                            result_data.append({
                                "url": current_url,
                                "title": title,
                                "file_path": str(temp_file)
                            })
                        elif action == "html":
                            html = await page.content()
                            result_data.append({
                                "url": current_url,
                                "title": title,
                                "html": html
                            })
            
            if result_data:
                return {"success": True, "pages": result_data}
            else:
                return {"success": False, "error": "No matching page found"}
    
    except Exception as e:
        return {"success": False, "error": f"CDP connection failed: {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" | "chrome-cdp",
        "cdp_port": 9222  // 用于 chrome-cdp 后端，连接到已打开的 Chrome
    }
    """
    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")
    cdp_port = int(data.get("cdp_port", 9222))
    url_hint = data.get("url_hint", "")
    
    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:16025")
    app.run(host='0.0.0.0', port=16025, debug=True)