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

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

app = FastAPI()

API_KEY = os.getenv('API_KEY')
if not API_KEY:
    raise ValueError("API_KEY environment variable must be set")

browser = None

async def get_browser():
    global browser
    if browser is None or not browser.isConnected():
        logger.info("Launching new browser instance")
        browser = await launch(
            headless=True,
            executablePath='/usr/bin/google-chrome',
            args=[
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-gpu',
                '--single-process',
                '--disable-software-rasterizer',
                '--js-flags=--expose-gc',
                '--memory-pressure-off',
            ]
        )
    return browser

async def cleanup_resources(page=None):
    try:
        if page:
            try:
                await page.evaluate("window.gc()")  # Trigger JS garbage collection
            except:
                pass  # Ignore if GC is not exposed
            await page.close()
        gc.collect()
        process = psutil.Process(os.getpid())
        logger.info(f"Memory usage: {process.memory_info().rss / 1024 / 1024:.2f} MB")
    except Exception as e:
        logger.error(f"Cleanup error: {str(e)}")

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

    page = None
    try:
        decoded_url = unquote(url)
        logger.info(f"Visiting URL: {decoded_url}")
        browser_instance = await get_browser()
        context = await browser_instance.createIncognitoBrowserContext()
        page = await context.newPage()
        await page.setViewport({"width": 1280, "height": 800})
        await page.goto(decoded_url, waitUntil='networkidle2')
        content = await page.content()
        if background_tasks:
            background_tasks.add_task(cleanup_resources, page)
        else:
            await cleanup_resources(page)
        return {"status": "success", "content": content}
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        if page:
            await cleanup_resources(page)
        raise HTTPException(status_code=500, detail=str(e))

@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Shutting down, cleaning up browser instance")
    global browser
    if browser:
        await browser.close()
        browser = None