#!/usr/bin/env python3
"""
Test script for verifying browser installation and accessibility in a Hugging Face Space.
This script attempts to:
1. Launch a browser using Playwright
2. Navigate to a basic website
3. Take a screenshot
4. Close the browser
"""

import os
import sys
import time
import argparse
from pathlib import Path
import logging
import subprocess

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

def check_system_dependencies():
    """Check for and install missing system dependencies."""
    if not sys.platform.startswith('linux'):
        logger.info("Not on Linux, skipping dependency check")
        return

    logger.info("Checking for system dependencies")
    try:
        # Check if we have a browser binary
        chrome_path = "/home/user/.cache/ms-playwright/chromium-1105/chrome-linux/chrome"
        if os.path.exists(chrome_path):
            # Try to check shared library dependencies using ldd
            logger.info(f"Checking dependencies for: {chrome_path}")
            result = subprocess.run(f"ldd {chrome_path}", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output = result.stdout.decode('utf-8', errors='ignore')
            
            if "not found" in output:
                logger.warning(f"Missing dependencies detected:\n{output}")
                
                # Try to install dependencies
                try:
                    logger.info("Installing Chrome dependencies")
                    pkg_cmd = """apt-get update && apt-get install -y --no-install-recommends \
                            libnss3 libnspr4 libasound2 libatk1.0-0 libatk-bridge2.0-0 \
                            libcups2 libdbus-1-3 libdrm2 libgbm1 libgtk-3-0 libxkbcommon0 \
                            libxcomposite1 libxdamage1 libxfixes3 libxrandr2 libgbm-dev \
                            libxshmfence1 libgles2 libegl1 xvfb fonts-liberation"""
                    
                    # Check if sudo is available
                    sudo_check = subprocess.run("which sudo", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                    if sudo_check.returncode == 0:
                        logger.info("Using sudo to install dependencies")
                        os.system(f"sudo {pkg_cmd}")
                    else:
                        logger.info("Trying direct install (no sudo)")
                        os.system(pkg_cmd)
                        
                    logger.info("Dependencies installed")
                except Exception as e:
                    logger.error(f"Failed to install dependencies: {e}")
    except Exception as e:
        logger.warning(f"Error checking dependencies: {e}")

def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Browser test script")
    parser.add_argument("--no-headless", action="store_true", help="Run in non-headless mode")
    args = parser.parse_args()
    
    try:
        logger.info("Starting browser test")
        from playwright.sync_api import sync_playwright
        
        logger.info("Imported playwright successfully")
        
        # Check system dependencies
        check_system_dependencies()
        
        # Check if we're in a Hugging Face Space
        is_hf_space = os.environ.get("SPACE_ID") is not None
        logger.info(f"Running in HF Space environment: {is_hf_space}")
        logger.info(f"Platform: {sys.platform}")
        logger.info(f"Headless mode: {not args.no_headless}")
        
        # Create output directory
        output_dir = Path("./screenshots")
        output_dir.mkdir(exist_ok=True)
        
        logger.info("Launching playwright")
        with sync_playwright() as playwright:
            browser_options = {
                "headless": not args.no_headless,
                "args": [
                    "--no-sandbox",
                    "--disable-gpu",
                    "--disable-dev-shm-usage",
                ]
            }
            
            # Add Linux-specific arguments
            if sys.platform.startswith('linux'):
                browser_options["args"].extend([
                    "--disable-setuid-sandbox",
                    "--single-process",
                ])
            
            # Linux-specific browser path handling
            if sys.platform.startswith('linux'):
                # Check for an existing browser config
                browser_config_path = Path("browser_config/system_browser.py")
                if browser_config_path.exists():
                    logger.info("Found system browser config file")
                    sys.path.append(str(Path("browser_config").absolute()))
                    try:
                        from system_browser import SYSTEM_BROWSER_PATH
                        if os.path.exists(SYSTEM_BROWSER_PATH) and os.access(SYSTEM_BROWSER_PATH, os.X_OK):
                            logger.info(f"Using system browser at {SYSTEM_BROWSER_PATH}")
                            browser_options["executable_path"] = SYSTEM_BROWSER_PATH
                    except ImportError:
                        logger.warning("Could not import system browser config")
                
                # Check for browser in cache
                chrome_path = "/home/user/.cache/ms-playwright/chromium-1105/chrome-linux/chrome"
                if os.path.exists(chrome_path) and os.access(chrome_path, os.X_OK):
                    logger.info(f"Found browser in cache at {chrome_path}")
                    browser_options["executable_path"] = chrome_path
            
            logger.info("Launching browser with options:")
            for key, value in browser_options.items():
                logger.info(f"  {key}: {value}")
                
            browser = playwright.chromium.launch(**browser_options)
            context = browser.new_context()
            page = context.new_page()
            
            logger.info("Browser launched successfully!")
            
            try:
                logger.info("Navigating to example.com")
                page.goto("https://example.com")
                
                logger.info("Page loaded, taking screenshot")
                screenshot_path = output_dir / "test_screenshot.png"
                page.screenshot(path=str(screenshot_path))
                
                logger.info(f"Screenshot saved to {screenshot_path}")
                logger.info(f"Page title: {page.title()}")
                
                # Get some page information
                dimensions = page.evaluate("""() => {
                    return {
                        width: window.innerWidth,
                        height: window.innerHeight,
                        devicePixelRatio: window.devicePixelRatio
                    }
                }""")
                logger.info(f"Window dimensions: {dimensions}")
                
            except Exception as e:
                logger.error(f"Error during page navigation: {e}")
            finally:
                logger.info("Closing browser")
                context.close()
                browser.close()
                
        logger.info("Test completed successfully!")
        return 0
        
    except Exception as e:
        logger.error(f"Test failed with error: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return 1

if __name__ == "__main__":
    sys.exit(main())