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

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")

# Define custom user agent
CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36'

async def wait_for_network_idle(page):
    """Wait until no network requests are in flight"""
    await page.waitForNetworkIdle(idleTime=500, timeout=30000)

@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")

    browser = None
    try:
        # Decode URL if it's encoded
        decoded_url = unquote(url)
        print(decoded_url)

        # Launch browser using installed Chrome
        browser = await launch(
            headless=True,
            executablePath='/usr/bin/google-chrome',
            args=['--no-sandbox', '--disable-setuid-sandbox']
        )

        # Create new page
        page = await browser.newPage()

        # Set custom user agent
        await page.setUserAgent(CUSTOM_USER_AGENT)

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

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

        # Close page
        await page.close()

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

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        # Ensure browser is closed even if an error occurs
        if browser:
            await browser.close()

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