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

import asyncio
from flask import Flask, request, jsonify
from proxy_lite import Runner, RunnerConfig
import os
import logging
from datetime import datetime
from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeoutError

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

app = Flask(__name__)

_runner = None 

async def perform_hardcoded_salesforce_login_and_get_cookies(username, password, login_url, target_url):
    logger.info("Attempting hardcoded Salesforce login with Playwright to obtain cookies...")
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True, args=["--no-sandbox", "--disable-setuid-sandbox"])
        context = await browser.new_context()
        page = await context.new_page()

        try:
            await page.goto(login_url, wait_until="domcontentloaded", timeout=60000)
            logger.info(f"Playwright: Navigated to Salesforce login page: {page.url}")

            await page.fill("#username", username)
            await page.fill("#password", password)
            await page.click("#Login")
            logger.info("Playwright: Filled credentials and clicked Login. Waiting for post-login state...")

            try:
                await page.wait_for_url(lambda url: "login.salesforce.com" not in url and "unauthorized" not in url.lower(), timeout=60000)
                logger.info(f"Playwright: Successfully redirected from login page. Current URL: {page.url}")
                await page.wait_for_selector('button[title="App Launcher"]', timeout=30000)
                logger.info("Playwright: Main Salesforce Lightning UI (e.g., App Launcher) detected after login.")

            except PlaywrightTimeoutError:
                logger.error(f"Playwright: Did not detect main UI or expected URL change within timeout after login. Current URL: {page.url}. Login might have failed or stuck on a redirect loop.")
                raise Exception("Salesforce login redirection failed or main UI not detected.")
            
            logger.info(f"Playwright: Navigating to target URL: {target_url} to ensure all relevant cookies are captured.")
            await page.goto(target_url, wait_until="domcontentloaded", timeout=60000)

            try:
                # Wait for generic Salesforce setup page elements to load
                await page.wait_for_selector('.setupPage, .slds-page-header, .slds-card, [data-aura-class*="setup"], .forcePageBlockSectionView', timeout=30000)
                logger.info("Playwright: Detected Salesforce setup page elements loaded successfully.")
            except PlaywrightTimeoutError:
                logger.warning("Playwright: Specific setup page elements not found. Trying generic page load check...")
                try:
                    # Fallback: wait for page to reach network idle state
                    await page.wait_for_load_state("networkidle", timeout=10000)
                    logger.info("Playwright: Page reached network idle state - proceeding with task.")
                except PlaywrightTimeoutError:
                    logger.info("Playwright: Page load validation timed out, but continuing as page may still be functional.")

            await asyncio.sleep(2)
            logger.info(f"Playwright: Successfully navigated to and confirmed content on {page.url}")

            cookies = await context.cookies()
            logger.info(f"Playwright: Extracted {len(cookies)} cookies after successful login and navigation.")
            return cookies

        except PlaywrightTimeoutError as e:
            logger.error(f"Playwright login/navigation failed (Timeout): {e}. Current URL: {page.url}")
            raise 
        except Exception as e:
            logger.error(f"Playwright login/navigation failed (General Error): {e}. Current URL: {page.url}")
            raise
        finally:
            if browser:
                await browser.close()


async def initialize_runner_with_cookies(cookies: list, target_url: str):
    global _runner
    logger.info("Initializing Proxy-lite Runner with provided cookies...")
    
    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_dict = {
        "environment": {
            "name": "webbrowser",
            "homepage": "about:blank",  # Safe startup, we'll open new tab programmatically
            "headless": True, 
            "launch_args": ["--no-sandbox", "--disable-setuid-sandbox"],
            "screenshot_delay": 0.5, 
            "include_html": True,
            "include_poi_text": True,
            "record_pois": True, 
            "viewport_width": 1280,
            "viewport_height": 720,
            "browserbase_timeout": 7200,
            "keep_original_image": False,
            "no_pois_in_image": False,
            "initial_cookies": cookies 
        },
        "solver": {
            "name": "simple",
            "agent": {
                "name": "proxy_lite", # Corrected as per previous error
                "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, 
                    "http_timeout": 50.0,
                    "http_concurrent_connections": 50,
                },
                "history_messages_limit": {
                    "screenshot": 1
                },
                "history_messages_include": None,
            }
        },
        "environment_timeout": 1800.0,
        "action_timeout": 1800.0,
        "task_timeout": 18000.0,
        "max_steps": 150,
        "logger_level": "DEBUG", 
        "save_every_step": True,
        "detailed_logger_name": False
    }
    config = RunnerConfig.from_dict(config_dict)

    logger.info(f"DEBUG: app.py - Initializing Proxy-lite Runner with config (cookies to be injected).")
    _runner = Runner(config=config)
    logger.info("Proxy-lite Runner initialized successfully with injected cookies.")
    return _runner


@app.route('/run_proxy_task', methods=['POST'])
async def run_proxy_task_endpoint():
    data = request.json
    request_task_instruction = data.get('task')
    target_url = data.get('url')

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

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

    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

    salesforce_login_url = "https://login.salesforce.com/"

    try:
        logger.info("Executing hardcoded login via Playwright to get session cookies...")
        session_cookies = await perform_hardcoded_salesforce_login_and_get_cookies(
            salesforce_username, salesforce_password, salesforce_login_url, target_url
        )
        logger.info(f"Successfully obtained {len(session_cookies)} cookies. These will be injected into the agent's browser.")

        runner = await initialize_runner_with_cookies(session_cookies, target_url)
        logger.info("Proxy-lite Runner initialized with pre-set cookies.")

        logger.info("Agent will use mandatory new tab tool to bypass loading issues.")
        
        # MANDATORY new tab navigation task - this is critical to avoid loading issues
        agent_task = f"""
CRITICAL FIRST STEP - MANDATORY: 
Your VERY FIRST action must be to use the open_new_tab_and_go_to tool to navigate to {target_url}

DO NOT skip this step. DO NOT use goto. You MUST use: open_new_tab_and_go_to(url='{target_url}')

This is necessary because direct navigation to this URL gets stuck loading. The new tab approach bypasses this issue.

STEP 1: Use open_new_tab_and_go_to(url='{target_url}')
STEP 2: Wait for the page to be fully loaded (no loading spinners visible)
STEP 3: {request_task_instruction}

Report success/failure for each step.
        """
        
        logger.info("Executing agent task with mandatory new tab navigation...")
        result = await runner.run(task=agent_task)

        # Extract the actual result value from the Run object
        if hasattr(result, 'value') and result.value:
            task_result = str(result.value)
        elif hasattr(result, 'result') and result.result:
            task_result = str(result.result)
        else:
            task_result = str(result)
            
        logger.info(f"Proxy-lite task completed. Output (truncated for log): {task_result[:500]}...")
        
        # Structure response for LWC integration
        response = {
            "status": "success",
            "message": "Task completed successfully",
            "data": {
                "task_result": task_result,
                "steps_completed": [
                    "Hardcoded Salesforce login completed",
                    "Browser session initialized with cookies",
                    "New tab navigation executed",
                    "Target Salesforce setup page accessed",
                    "Task execution completed successfully"
                ],
                "environment": {
                    "target_url": target_url,
                    "cookies_count": len(session_cookies),
                    "navigation_method": "new_tab_bypass"
                }
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        
        return jsonify(response)
        
    except PlaywrightTimeoutError as e:
        logger.exception(f"Playwright timeout during login/navigation: {e}")
        error_response = {
            "status": "error",
            "error_type": "navigation_timeout",
            "message": "Page loading timed out during login or navigation",
            "data": {
                "error_details": str(e),
                "suggested_action": "Retry the request - network issues may be temporary",
                "steps_completed": ["Login attempted", "Navigation failed due to timeout"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 500
        
    except ValueError as e:
        logger.exception(f"Configuration error: {e}")
        error_response = {
            "status": "error", 
            "error_type": "configuration_error",
            "message": "System configuration issue",
            "data": {
                "error_details": str(e),
                "suggested_action": "Check environment variables and system configuration",
                "steps_completed": ["Configuration validation failed"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 500
        
    except Exception as e:
        logger.exception(f"Unexpected error processing Salesforce task: {e}")
        error_response = {
            "status": "error",
            "error_type": "unexpected_error", 
            "message": "An unexpected error occurred during task execution",
            "data": {
                "error_details": str(e),
                "error_class": type(e).__name__,
                "suggested_action": "Check logs for detailed error information and retry",
                "steps_completed": ["Login attempted", "Error occurred during execution"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 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."

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint for monitoring and debugging"""
    logger.info("Health check endpoint accessed.")
    
    # Check environment variables
    env_status = {
        "HF_API_TOKEN": "✓" if os.environ.get("HF_API_TOKEN") else "✗",
        "SALESFORCE_USERNAME": "✓" if os.environ.get("SALESFORCE_USERNAME") else "✗", 
        "SALESFORCE_PASSWORD": "✓" if os.environ.get("SALESFORCE_PASSWORD") else "✗"
    }
    
    health_response = {
        "status": "healthy",
        "message": "Proxy-lite API is running",
        "environment_variables": env_status,
        "endpoints": {
            "POST /run_proxy_task": "Execute Salesforce automation tasks (requires 'task' and 'url' parameters)",
            "GET /health": "Health check and status",
            "GET /": "API information"
        },
        "supported_pages": [
            "Warranty Lifecycle Management",
            "Account Forecasting Settings", 
            "Sales Agreements",
            "Account Manager Targets",
            "Any Salesforce Setup page"
        ],
        "timestamp": datetime.now().isoformat()
    }
    
    return jsonify(health_response)

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.")
    logger.info("Starting Flask development server on 0.0.0.0:7860...")
    app.run(host='0.0.0.0', port=7860, debug=True)