from fastapi import FastAPI, HTTPException, Header
from pyppeteer import launch
import os
import asyncio
from typing import Optional
from urllib.parse import unquote
import logging
import gc

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Get API key from environment variable
API_KEY = os.getenv('API_KEY')
if not API_KEY:
    raise ValueError("API_KEY environment variable must be set")

# Browser pool
browser = None

# Browser configuration
BROWSER_ARGS = [
    '--no-sandbox',
    '--disable-setuid-sandbox',
    '--disable-dev-shm-usage',  # Overcome limited resource problems
    '--disable-gpu',            # Disable GPU hardware acceleration
    '--disable-extensions',     # Disable extensions
    '--disable-audio-output',   # Disable audio
    '--disable-background-networking',
    '--disable-background-timer-throttling',
    '--disable-backgrounding-occluded-windows',
    '--disable-breakpad',
    '--disable-component-extensions-with-background-pages',
    '--disable-features=TranslateUI,BlinkGenPropertyTrees',
    '--disable-ipc-flooding-protection',
    '--disable-renderer-backgrounding',
    '--enable-features=NetworkService,NetworkServiceInProcess',
    '--mute-audio',
    '--memory-pressure-off',
]

# Browser launch options
BROWSER_OPTIONS = {
    'headless': True,
    'executablePath': '/usr/bin/google-chrome',
    'args': BROWSER_ARGS,
    'handleSIGINT': False,
    'handleSIGTERM': False,
    'handleSIGHUP': False,
}

async def get_browser():
    """Get or create a browser instance"""
    global browser
    if browser is None or not browser.isConnected():
        logger.info("Launching new browser instance")
        browser = await launch(**BROWSER_OPTIONS)
    return browser

async def wait_for_network_idle(page):
    """Wait until no network requests are in flight"""
    try:
        await page.waitForNetworkIdle(idleTime=500, timeout=30000)
    except Exception as e:
        logger.warning(f"Network idle timeout: {str(e)}")

@app.on_event("startup")
async def startup_event():
    """Initialize browser on startup"""
    await get_browser()

@app.on_event("shutdown")
async def shutdown_event():
    """Close browser on shutdown"""
    global browser
    if browser:
        logger.info("Closing browser on shutdown")
        await browser.close()
        browser = None

@app.head("/")
async def health_check():
    return {"status": "ok"}

@app.get("/")
async def visit_url(url: str, x_api_key: Optional[str] = Header(None)):
    # Validate API key
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

    page = None
    try:
        # Decode URL if it's encoded
        decoded_url = unquote(url)
        logger.info(f"Visiting URL: {decoded_url}")

        # Get browser instance
        browser_instance = await get_browser()

        # Create new page with timeout
        page = await browser_instance.newPage()

        # Set page timeout and resource limits
        await page.setDefaultNavigationTimeout(60000)  # 60 seconds timeout
        await page.setRequestInterception(True)

        # Block unnecessary resources
        async def intercept_request(request):
            if request.resourceType in ['image', 'media', 'font']:
                await request.abort()
            else:
                await request.continue_()

        page.on('request', intercept_request)

        # Navigate to URL and wait for network idle
        await page.goto(decoded_url, waitUntil='networkidle2', timeout=60000)

        # Get page content
        content = await page.content()

        # Close page to free up resources
        await page.close()

        # Force garbage collection
        gc.collect()

        return {"status": "success", "content": content}

    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        # Ensure browser is still running
        if browser is None or not browser.isConnected():
            await get_browser()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        # Always ensure page is closed
        if page:
            try:
                await page.close()
            except:
                pass

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)