from fastapi import FastAPI, Query
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import base64
from playwright.async_api import async_playwright

app = FastAPI()

# Optional: Allow access from any frontend (useful for Hugging Face Gradio frontends)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/screenshot")
async def screenshot(url: str = Query(..., description="Web page URL to capture")):
    try:
        async with async_playwright() as p:
            browser = await p.chromium.launch()
            page = await browser.new_page()
            await page.goto(url, timeout=30000)  # 30s timeout
            buffer = await page.screenshot(full_page=True)
            await browser.close()

        b64_image = base64.b64encode(buffer).decode("utf-8")
        data_url = f"data:image/png;base64,{b64_image}"
        return JSONResponse(content={"image_base64": data_url})

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})