"""
LANDRUN + BROWSER-USE + CHROMIUM MERGED SYSTEM
==============================================
Kernel-level code execution sandbox with AI-powered browser automation

Features:
1. Landrun: Go-based Linux Landlock kernel security sandbox
2. Browser-Use: AI agent for intelligent browser automation
3. Chromium: Playwright browser for UI testing
4. FastAPI: Modern async web framework

Endpoints:
- POST /execute - Execute code in Landrun sandbox
- GET /preview/{id} - Get live HTML preview
- POST /browser/test - Test UI with Playwright
- POST /browser/agent - AI agent automated browsing
- POST /browser/execute_and_agent - One-shot: Execute + AI Agent
"""

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Any
import subprocess
import tempfile
import os
import base64
import uuid
from datetime import datetime, timedelta
import asyncio
import json

# Playwright for direct browser control
from playwright.async_api import async_playwright

# Browser-Use for AI agent automation
try:
    from browser_use import Agent
    from langchain_openai import ChatOpenAI
    BROWSER_USE_AVAILABLE = True
except ImportError:
    BROWSER_USE_AVAILABLE = False
    print("⚠️  Browser-Use not available - AI agent features disabled")

app = FastAPI(
    title="Landrun + Browser-Use + Chromium",
    description="Kernel-level sandbox with AI browser automation",
    version="2.0.0"
)

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Storage
PREVIEW_STORAGE = {}
PREVIEW_EXPIRY = timedelta(hours=1)

# ============================================================================
# PYDANTIC MODELS
# ============================================================================

class CodeExecutionRequest(BaseModel):
    language: str = Field(..., description="Language: python, javascript, react, html")
    code: str = Field(..., description="Source code to execute")

class BrowserAction(BaseModel):
    type: str = Field(..., description="Action type: click, type, get_text, wait, screenshot")
    selector: Optional[str] = Field(None, description="CSS selector for element")
    text: Optional[str] = Field(None, description="Text to type (for type action)")
    timeout: Optional[int] = Field(5000, description="Timeout in milliseconds")

class BrowserTestRequest(BaseModel):
    preview_url: str = Field(..., description="Preview URL to test")
    actions: List[BrowserAction] = Field(..., description="List of browser actions")

class BrowserAgentRequest(BaseModel):
    task: str = Field(..., description="Natural language task for AI agent")
    url: Optional[str] = Field(None, description="Starting URL (optional)")
    max_steps: Optional[int] = Field(10, description="Maximum number of steps")

class ExecuteAndAgentRequest(BaseModel):
    language: str = Field(..., description="Language: python, javascript, react, html")
    code: str = Field(..., description="Source code to execute")
    agent_task: str = Field(..., description="AI agent task to perform on preview")
    max_steps: Optional[int] = Field(10, description="Maximum agent steps")

# ============================================================================
# LANDRUN CODE EXECUTION
# ============================================================================

def execute_with_landrun(language: str, code: str) -> dict:
    """Execute code using landrun kernel-level sandboxing"""
    
    configs = {
        "python": {
            "ext": ".py",
            "cmd": ["python3"],
            "allowed_paths": ["/usr/lib/python3*", "/usr/local/lib/python3*"],
        },
        "javascript": {
            "ext": ".js",
            "cmd": ["node"],
            "allowed_paths": ["/usr/lib/node_modules", "/usr/local/lib/node_modules"],
        },
        "html": {
            "ext": ".html",
            "cmd": None,
            "allowed_paths": [],
        },
        "react": {
            "ext": ".jsx",
            "cmd": ["node"],
            "allowed_paths": ["/usr/lib/node_modules", "/usr/local/lib/node_modules"],
        }
    }
    
    config = configs.get(language.lower())
    if not config:
        return {"error": f"Unsupported language: {language}"}
    
    try:
        os.makedirs('/tmp/sandbox', exist_ok=True)
        
        with tempfile.NamedTemporaryFile(mode='w', suffix=config['ext'], delete=False, dir='/tmp/sandbox') as f:
            f.write(code)
            temp_file = f.name
        
        # HTML - return directly
        if language.lower() == "html":
            with open(temp_file, 'r') as f:
                html_content = f.read()
            os.unlink(temp_file)
            return {
                "output": "HTML rendered successfully",
                "preview": base64.b64encode(html_content.encode()).decode(),
                "exit_code": 0
            }
        
        # React - wrap and transpile
        if language.lower() == "react":
            react_wrapper = f"""
import React from 'react';
import {{ createRoot }} from 'react-dom/client';

{code}

const root = createRoot(document.getElementById('root'));
root.render(<App />);
"""
            html_template = """
<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <script crossorigin src="https://unpkg.com/react@18/umd/react.production.min.js"></script>
    <script crossorigin src="https://unpkg.com/react-dom@18/umd/react-dom.production.min.js"></script>
    <script src="https://unpkg.com/@babel/standalone/babel.min.js"></script>
</head>
<body>
    <div id="root"></div>
    <script type="text/babel">
        {CODE}
    </script>
</body>
</html>
""".replace("{CODE}", code)
            
            os.unlink(temp_file)
            return {
                "output": "React component compiled",
                "preview": base64.b64encode(html_template.encode()).decode(),
                "exit_code": 0
            }
        
        # Build landrun command
        landrun_cmd = [
            "/usr/local/bin/landrun",
            "--ldd",
            "--add-exec",
            "--ro", "/usr",
            "--ro", "/lib",
            "--ro", "/lib64",
            "--ro", "/etc",
            "--rw", "/tmp/sandbox",
            "--ro", temp_file,
            "--connect-tcp", "80,443",
            "--log-level", "error",
        ]
        
        for path in config['allowed_paths']:
            landrun_cmd.extend(["--ro", path])
        
        landrun_cmd.append("--")
        landrun_cmd.extend(config['cmd'])
        landrun_cmd.append(temp_file)
        
        result = subprocess.run(
            landrun_cmd,
            capture_output=True,
            text=True,
            timeout=10
        )
        
        os.unlink(temp_file)
        
        return {
            "output": result.stdout + result.stderr,
            "exit_code": result.returncode,
            "security": "Landrun kernel-level isolation active"
        }
        
    except subprocess.TimeoutExpired:
        return {"error": "Execution timeout (10s)", "exit_code": -1}
    except Exception as e:
        return {"error": str(e), "exit_code": -1}

# ============================================================================
# PLAYWRIGHT BROWSER AUTOMATION
# ============================================================================

async def run_browser_test(preview_url: str, actions: List[BrowserAction]) -> dict:
    """Run Playwright browser test with actions"""
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        page = await browser.new_page()
        
        results = []
        screenshot_initial = None
        screenshot_final = None
        
        try:
            # Navigate to preview
            await page.goto(preview_url, wait_until="networkidle", timeout=10000)
            await page.wait_for_timeout(1000)
            
            # Initial screenshot
            screenshot_initial = base64.b64encode(await page.screenshot()).decode()
            
            # Execute actions
            for action in actions:
                try:
                    if action.type == "click":
                        await page.click(action.selector, timeout=action.timeout)
                        results.append({"action": "click", "selector": action.selector, "status": "success"})
                    
                    elif action.type == "type":
                        await page.fill(action.selector, action.text, timeout=action.timeout)
                        results.append({"action": "type", "selector": action.selector, "text": action.text, "status": "success"})
                    
                    elif action.type == "get_text":
                        text = await page.text_content(action.selector, timeout=action.timeout)
                        results.append({"action": "get_text", "selector": action.selector, "text": text, "status": "success"})
                    
                    elif action.type == "wait":
                        await page.wait_for_selector(action.selector, timeout=action.timeout)
                        results.append({"action": "wait", "selector": action.selector, "status": "success"})
                    
                    elif action.type == "screenshot":
                        screenshot_final = base64.b64encode(await page.screenshot()).decode()
                        results.append({"action": "screenshot", "status": "success"})
                    
                    await page.wait_for_timeout(500)
                    
                except Exception as e:
                    results.append({"action": action.type, "selector": action.selector, "status": "error", "error": str(e)})
            
            # Final screenshot if not taken
            if not screenshot_final:
                screenshot_final = base64.b64encode(await page.screenshot()).decode()
            
        finally:
            await browser.close()
        
        return {
            "test_results": results,
            "screenshot_initial": screenshot_initial,
            "screenshot_final": screenshot_final
        }

# ============================================================================
# BROWSER-USE AI AGENT
# ============================================================================

async def run_ai_agent(task: str, url: Optional[str] = None, max_steps: int = 10) -> dict:
    """Run Browser-Use AI agent with natural language task"""
    
    if not BROWSER_USE_AVAILABLE:
        return {
            "status": "error",
            "error": "Browser-Use not installed. Install with: pip install browser-use langchain-openai"
        }
    
    try:
        # Check for OpenAI API key
        if not os.getenv("OPENAI_API_KEY"):
            return {
                "status": "error",
                "error": "OPENAI_API_KEY environment variable not set"
            }
        
        # Create AI agent
        llm = ChatOpenAI(model="gpt-4o")
        
        agent = Agent(
            task=task,
            llm=llm,
            max_steps=max_steps
        )
        
        # Run agent
        result = await agent.run()
        
        return {
            "status": "success",
            "task": task,
            "result": str(result),
            "steps_taken": len(result.history()) if hasattr(result, 'history') else 0
        }
        
    except Exception as e:
        return {
            "status": "error",
            "error": str(e)
        }

# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.post("/execute")
async def execute_code(request: CodeExecutionRequest):
    """Execute code in Landrun sandbox"""
    
    result = execute_with_landrun(request.language, request.code)
    
    # Store preview if available
    if "preview" in result and not "error" in result:
        preview_id = str(uuid.uuid4())
        preview_html = base64.b64decode(result["preview"]).decode()
        
        PREVIEW_STORAGE[preview_id] = {
            "html": preview_html,
            "created": datetime.now()
        }
        
        result["preview_url"] = f"/preview/{preview_id}"
        del result["preview"]
    
    return result

@app.get("/preview/{preview_id}")
async def get_preview(preview_id: str):
    """Get live HTML preview"""
    
    if preview_id not in PREVIEW_STORAGE:
        raise HTTPException(status_code=404, detail="Preview not found or expired")
    
    # Check expiry
    preview_data = PREVIEW_STORAGE[preview_id]
    if datetime.now() - preview_data["created"] > PREVIEW_EXPIRY:
        del PREVIEW_STORAGE[preview_id]
        raise HTTPException(status_code=410, detail="Preview expired")
    
    return HTMLResponse(content=preview_data["html"])

@app.post("/browser/test")
async def browser_test(request: BrowserTestRequest):
    """Test UI with Playwright browser automation"""
    
    # Build full URL if relative
    if request.preview_url.startswith("/preview/"):
        base_url = os.getenv("SPACE_HOST", "http://localhost:7860")
        full_url = f"{base_url}{request.preview_url}"
    else:
        full_url = request.preview_url
    
    result = await run_browser_test(full_url, request.actions)
    
    return {
        "status": "success",
        "url_tested": full_url,
        **result
    }

@app.post("/browser/agent")
async def browser_agent(request: BrowserAgentRequest):
    """Run AI agent for automated browsing"""
    
    result = await run_ai_agent(
        task=request.task,
        url=request.url,
        max_steps=request.max_steps
    )
    
    return result

@app.post("/browser/execute_and_test")
async def execute_and_test(request: CodeExecutionRequest):
    """Execute code and test with Playwright (existing endpoint for compatibility)"""
    
    # Execute code
    exec_result = execute_with_landrun(request.language, request.code)
    
    if "error" in exec_result:
        return {
            "status": "error",
            "execution": exec_result
        }
    
    # Store preview
    if "preview" in exec_result:
        preview_id = str(uuid.uuid4())
        preview_html = base64.b64decode(exec_result["preview"]).decode()
        
        PREVIEW_STORAGE[preview_id] = {
            "html": preview_html,
            "created": datetime.now()
        }
        
        preview_url = f"/preview/{preview_id}"
        exec_result["preview_url"] = preview_url
        del exec_result["preview"]
    else:
        return {
            "status": "error",
            "error": "No preview generated"
        }
    
    return {
        "status": "success",
        "execution": exec_result
    }

@app.post("/browser/execute_and_agent")
async def execute_and_agent(request: ExecuteAndAgentRequest):
    """ONE-SHOT: Execute code + Run AI agent on preview"""
    
    # Execute code
    exec_result = execute_with_landrun(request.language, request.code)
    
    if "error" in exec_result:
        return {
            "status": "error",
            "execution": exec_result
        }
    
    # Store preview
    if "preview" in exec_result:
        preview_id = str(uuid.uuid4())
        preview_html = base64.b64decode(exec_result["preview"]).decode()
        
        PREVIEW_STORAGE[preview_id] = {
            "html": preview_html,
            "created": datetime.now()
        }
        
        preview_url = f"/preview/{preview_id}"
        base_url = os.getenv("SPACE_HOST", "http://localhost:7860")
        full_preview_url = f"{base_url}{preview_url}"
        
        # Run AI agent on the preview
        agent_result = await run_ai_agent(
            task=f"{request.agent_task}. Start at URL: {full_preview_url}",
            url=full_preview_url,
            max_steps=request.max_steps
        )
        
        return {
            "status": "success",
            "execution": {
                **exec_result,
                "preview_url": preview_url
            },
            "agent": agent_result
        }
    else:
        return {
            "status": "error",
            "error": "No preview generated for AI agent"
        }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "landrun": "active",
        "browser": "playwright-chromium",
        "browser_use": "available" if BROWSER_USE_AVAILABLE else "not installed",
        "ai_agent": "enabled" if (BROWSER_USE_AVAILABLE and os.getenv("OPENAI_API_KEY")) else "disabled"
    }

@app.get("/")
async def root():
    """Root endpoint with API documentation"""
    return {
        "service": "Landrun + Browser-Use + Chromium",
        "version": "2.0.0",
        "features": {
            "landrun": "Kernel-level code execution sandbox",
            "playwright": "Direct browser automation",
            "browser_use": "AI agent for intelligent browsing",
            "chromium": "Headless browser engine"
        },
        "endpoints": {
            "POST /execute": "Execute code in sandbox",
            "GET /preview/{id}": "Get live HTML preview",
            "POST /browser/test": "Test UI with Playwright",
            "POST /browser/agent": "Run AI agent task",
            "POST /browser/execute_and_test": "Execute + Playwright test",
            "POST /browser/execute_and_agent": "Execute + AI agent (ONE-SHOT)"
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)