src/agents/browser.py

"""
Browser management for the leaderboard agent.
"""
import os
import gc
import logging
from io import BytesIO
from time import sleep

import helium
from PIL import Image
from selenium import webdriver
from smolagents import CodeAgent
from smolagents.agents import ActionStep

# Configuration du logger
logger = logging.getLogger("leaderboard-parser")

# Global driver variable
driver = None


def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:
    """
    Save a screenshot of the current browser state in memory for the agent.
    This is used as a callback for the agent to visualize the page.
    The screenshot is only kept in memory and not saved to disk.
    """
    sleep(2.0)  # Increased to allow time for JavaScript animations
    current_step = memory_step.step_number
    if driver is not None:
        for previous_memory_step in agent.memory.steps:  # Remove previous screenshots from logs for lean processing
            if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= current_step - 2:
                previous_memory_step.observations_images = None
        
        # Capture screenshot for agent visualization only (not saved to disk)
        png_bytes = driver.get_screenshot_as_png()
        image = Image.open(BytesIO(png_bytes))
        print(f"Captured a browser screenshot for agent: {image.size} pixels")
        memory_step.observations_images = [image.copy()]  # Create a copy to ensure it persists, important!

    # Update observations with current URL
    url_info = f"Current url: {driver.current_url}"
    memory_step.observations = (
        url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
    )
    return


def initialize_driver():
    """
    Initialize the Selenium WebDriver.
    Returns a configured Chrome WebDriver instance.
    """
    global driver
    
    # Si le driver existe déjà, on le nettoie d'abord pour éviter les fuites mémoire
    if driver is not None:
        close_driver()
    
    print("Démarrage de l'initialisation du navigateur Chrome...")
    
    # Configuration minimale mais suffisante pour éviter les timeouts
    chrome_options = webdriver.ChromeOptions()
    # Options essentielles pour l'environnement conteneurisé
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    # Options pour améliorer les performances
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--disable-extensions")
    
    try:
        print("Tentative de démarrage de Chrome avec Helium...")
        # Utiliser les options minimales mais avec le mode headless
        driver = helium.start_chrome(headless=True, options=chrome_options)
        print("Chrome démarré avec succès!")
        
        # Informations sur le navigateur
        print(f"Version de Chrome: {driver.capabilities.get('browserVersion', 'Inconnue')}")
        print(f"Plateforme: {driver.capabilities.get('platformName', 'Inconnue')}")
        
        # Augmenter le délai d'attente pour le chargement des pages
        driver.set_page_load_timeout(60)  # Augmenté à 60 secondes
        # Augmenter le délai d'attente pour les scripts
        driver.set_script_timeout(60)  # Augmenté à 60 secondes
        
        return driver
    except Exception as e:
        print(f"ERREUR lors du démarrage de Chrome: {str(e)}")
        # Capturer la trace complète pour diagnostic
        import traceback
        print("Trace d'erreur complète:")
        traceback.print_exc()
        
        # Vérifier si Chrome est disponible
        try:
            import subprocess
            chrome_path = os.environ.get("CHROME_PATH", "/usr/bin/google-chrome-stable")
            chrome_version_cmd = f"{chrome_path} --version"
            version_output = subprocess.check_output(chrome_version_cmd, shell=True, stderr=subprocess.STDOUT).decode()
            print(f"Version de Chrome installée: {version_output.strip()}")
        except Exception as chrome_check_error:
            print(f"Impossible de vérifier la version de Chrome: {str(chrome_check_error)}")
        
        raise


def close_driver():
    """
    Close the browser and clean up resources.
    """
    global driver
    
    try:
        print("Fermeture du navigateur et nettoyage des ressources...")
        
        # Utiliser helium.kill_browser() pour fermer proprement le navigateur
        helium.kill_browser()
        
        # Libérer la référence
        driver = None
        
        # Forcer le garbage collector
        gc.collect()
        
        print("Navigateur fermé avec succès")
    except Exception as e:
        print(f"Error closing browser: {e}")


# Alias de close_driver pour compatibilité avec browser_utils.cleanup_browser
def cleanup_browser():
    """
    Alias de close_driver pour compatibilité avec l'API existante.
    """
    close_driver()