from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import asyncio
import os
import sys
from threading import Thread
import json
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import browser-use components and LLM libraries
try:
    from langchain_openai import ChatOpenAI
    from langchain_anthropic import ChatAnthropic
    from langchain_google_genai import ChatGoogleGenerativeAI
    from browser_use import Agent
except ImportError as e:
    print(f"Import error: {e}")
    print("Make sure to install browser-use, langchain-openai, langchain-anthropic, and langchain-google-genai")

app = Flask(__name__, static_folder="static")
CORS(app)  # Enable CORS for all routes

# Global variables for managing async operations
loop = None
thread = None

def start_background_loop(loop):
    """Start the asyncio event loop in a background thread"""
    asyncio.set_event_loop(loop)
    loop.run_forever()

def get_or_create_eventloop():
    """Get or create the asyncio event loop"""
    global loop, thread
    if loop is None:
        loop = asyncio.new_event_loop()
        thread = Thread(target=start_background_loop, args=(loop,))
        thread.daemon = True
        thread.start()
    return loop

async def run_browser_task(task, model="gpt-4o-mini"):
    """Run a browser automation task using browser-use"""
    try:
        # Initialize the LLM based on the model name
        if model.startswith("gpt"):
            llm = ChatOpenAI(model=model, temperature=0.1)
        elif model.startswith("claude"):
            llm = ChatAnthropic(model=model, temperature=0.1)
        elif model.startswith("gemini"):
            llm = ChatGoogleGenerativeAI(model=model, temperature=0.1)
        else:
            raise ValueError(f"Unsupported model: {model}")
        
        # Create the agent
        agent = Agent(task=task, llm=llm)
        
        # Run the task
        result = await agent.run()
        
        return {
            "success": True,
            "result": str(result),
            "task": task
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "task": task
        }

@app.route("/")
def home():
    """Serve the HTML interface"""
    return send_from_directory("static", "index.html")

@app.route("/api")
def api_docs():
    """API documentation endpoint"""
    return jsonify({
        "message": "Browser-Use Server API",
        "endpoints": {
            "/": "Web interface",
            "/api": "This API documentation",
            "/health": "Health check",
            "/run-task": "POST - Run a browser automation task",
            "/status": "Get server status"
        },
        "usage": {
            "run_task": {
                "method": "POST",
                "url": "/run-task",
                "body": {
                    "task": "Description of the task to perform",
                    "model": "gpt-4o-mini (optional, default)"
                },
                "example": {
                    "task": "Go to google.com and search for 'Python programming'",
                    "model": "gpt-4o-mini"
                }
            }
        }
    })

@app.route("/health")
def health():
    """Health check endpoint"""
    return jsonify({
        "status": "healthy",
        "service": "browser-use-server",
        "version": "1.0.0"
    })

@app.route("/status")
def status():
    """Get server status and configuration"""
    return jsonify({
        "status": "running",
        "environment": {
            "python_version": sys.version,
            "has_openai_key": bool(os.getenv("OPENAI_API_KEY")),
            "has_anthropic_key": bool(os.getenv("ANTHROPIC_API_KEY")),
            "has_google_api_key": bool(os.getenv("GOOGLE_API_KEY"))
        },
        "supported_models": [
            "gpt-4o",
            "gpt-4o-mini", 
            "gpt-4-turbo",
            "claude-3-sonnet-20240229",
            "claude-3-haiku-20240307",
            "gemini-2.0-flash",
            "gemini-2.5-flash",
            "gemini-2.0-flash-exp",
            "gemini-2.0-flash-lite"
        ]
    })

@app.route("/run-task", methods=["POST"])
def run_task():
    """Run a browser automation task"""
    try:
        # Get request data
        data = request.get_json()
        
        if not data or "task" not in data:
            return jsonify({
                "success": False,
                "error": "Missing 'task' in request body"
            }), 400
        
        task = data["task"]
        model = data.get("model", "gpt-4o-mini")
        
        # Validate API keys
        if not os.getenv("OPENAI_API_KEY") and not os.getenv("ANTHROPIC_API_KEY") and not os.getenv("GOOGLE_API_KEY"):
            return jsonify({
                "success": False,
                "error": "No API keys configured. Please set OPENAI_API_KEY, ANTHROPIC_API_KEY, or GOOGLE_API_KEY environment variable."
            }), 500
        
        # Get the event loop
        loop = get_or_create_eventloop()
        
        # Run the task asynchronously
        future = asyncio.run_coroutine_threadsafe(
            run_browser_task(task, model), 
            loop
        )
        
        # Wait for the result (with timeout)
        try:
            result = future.result(timeout=300)  # 5 minute timeout
            return jsonify(result)
        except asyncio.TimeoutError:
            return jsonify({
                "success": False,
                "error": "Task timed out after 5 minutes"
            }), 408
            
    except Exception as e:
        return jsonify({
            "success": False,
            "error": f"Server error: {str(e)}"
        }), 500

if __name__ == "__main__":
    print("Starting Browser-Use Server...")
    print("Make sure you have set your API keys:")
    print("- OPENAI_API_KEY for OpenAI models")
    print("- ANTHROPIC_API_KEY for Anthropic models")
    print("- GOOGLE_API_KEY for Google Gemini models")
    
    # Run the Flask app
    app.run(host="0.0.0.0", port=7860, debug=False)