import gevent.monkey
gevent.monkey.patch_all(asyncio=True) # Keep this at the very top

import asyncio # Keep this
from flask import Flask, request, jsonify
from proxy_lite import Runner, RunnerConfig
import os
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

_runner = None

async def initialize_runner():
    global _runner
    if _runner is None:
        logger.info("Initializing Proxy-lite Runner...")
        
        hf_api_token = os.environ.get("HF_API_TOKEN")
        if not hf_api_token:
            logger.error("HF_API_TOKEN environment variable not set. Cannot initialize Runner.")
            raise ValueError("HF_API_TOKEN environment variable not set. Please set it as a Space secret.")

        config = RunnerConfig.from_dict({
            "environment": {
                "name": "webbrowser",
                "homepage": "https://dwd000006jia1mae.lightning.force.com/lightning/setup/AccountForecastSettings/home",
                "headless": True,
                "launch_args": ["--no-sandbox", "--disable-setuid-sandbox"],
                "screenshot_delay": 3.0,
                "include_html": True,
                "include_poi_text": True,
            },
            "solver": {
                "name": "simple",
                "agent": {
                    "name": "proxy_lite",
                    "client": {
                        "name": "convergence",
                        "model_id": "convergence-ai/proxy-lite-3b",
                        "api_base": "https://convergence-ai-demo-api.hf.space/v1", 
                        "api_key": hf_api_token
                    }
                }
            },
            "environment_timeout": 1800.0,
            "action_timeout": 1800.0,
            "task_timeout": 18000.0,
            "max_steps": 150,
            "logger_level": "DEBUG",
        })

        logger.info(f"DEBUG: app.py - Initializing Runner with environment_timeout: {config.environment_timeout} seconds")
        logger.info(f"DEBUG: app.py - Full config used: {config.model_dump_json(indent=2)}")

        _runner = Runner(config=config)
        logger.info("Proxy-lite Runner initialized successfully.")
    return _runner

# --- MODIFIED run_async_task FUNCTION (SIMPLIFIED) ---
# This function is no longer needed in most cases with gevent.monkey.patch_all(asyncio=True)
# but if you must call async functions from sync context, you simply await them.
# However, you are already in an async function context within Flask routes when using Gunicorn/gevent.
# The Gunicorn worker itself implicitly runs an event loop.
# Let's remove the run_until_complete part.

# DELETED: def run_async_task(coro): ...

# --- END MODIFIED run_async_task FUNCTION ---


@app.route('/run_proxy_task', methods=['POST'])
async def run_proxy_task_endpoint(): # <--- MAKE THIS FUNCTION ASYNC
    data = request.json
    request_task_instruction = data.get('task')

    if not request_task_instruction:
        logger.warning("Received request without 'task' field. Returning 400.")
        return jsonify({"error": "No 'task' provided in request body"}), 400

    logger.info(f"Received user request task: '{request_task_instruction}'")

    salesforce_username = os.environ.get("SALESFORCE_USERNAME")
    salesforce_password = os.environ.get("SALESFORCE_PASSWORD")

    if not salesforce_username or not salesforce_password:
        logger.error("Salesforce credentials (SALESFORCE_USERNAME, SALESFORCE_PASSWORD) environment variables not set.")
        return jsonify({"error": "Salesforce credentials not configured. Please set SALESFORCE_USERNAME and SALESFORCE_PASSWORD as Space secrets."}), 500

    agent_task = (
        f"Log in to Salesforce. The username is '{salesforce_username}' and the password is '{salesforce_password}'. "
        f"After attempting to log in, observe the page carefully. "
        f"If the login was successful, the URL should change from the login page, and you should see elements indicating a logged-in state (e.g., a Salesforce navigation menu, a home screen, or a profile icon), rather than a login form or an error message. "
        f"If the login is successful, {request_task_instruction}. "
        f"Report the final status of the requested action and confirmation of successful login."
    )
    
    logger.info(f"Executing agent task: '{agent_task[:200]}...'")

    try:
        # Since run_proxy_task_endpoint is now async, you can directly await
        runner = await initialize_runner()
        result = await runner.run(agent_task) # <--- AWAIT DIRECTLY

        logger.info(f"Proxy-lite task completed. Output: {result[:200]}...")
        return jsonify({"output": result})
    except Exception as e:
        logger.exception(f"Error processing Salesforce task: {e}")
        # The RuntimeWarning: coroutine 'initialize_runner' was never awaited will disappear
        # because initialize_runner is now awaited.
        return jsonify({"error": f"An error occurred: {str(e)}. Check logs for details."}), 500

@app.route('/')
def root():
    logger.info("Root endpoint accessed.")
    return "Proxy-lite API is running. Send POST requests to /run_proxy_task with a 'task' in JSON body."

if __name__ == '__main__':
    if not os.environ.get("HF_API_TOKEN"):
        logger.error("HF_API_TOKEN environment variable is not set. Please set it for local testing.")
        exit(1)
    logger.info("Starting Flask development server on 0.0.0.0:7860...")
    app.run(host='0.0.0.0', port=7860, debug=True)