import streamlit as st
from playwright.async_api import async_playwright
import asyncio
from io import BytesIO

st.set_page_config(page_title="Webpage Screenshot", layout="wide")
st.title("All-Purpose Webpage Screenshot")

url = st.text_input(
    "Enter any webpage URL",
    value="https://example.com"
)

if not url.startswith("http"):
    st.warning("Enter a valid URL (http / https)")
    st.stop()


def get_screenshot(url: str):

    async def run():
        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=True,
                args=[
                    "--no-sandbox",
                    "--disable-blink-features=AutomationControlled"
                ]
            )

            context = await browser.new_context(
                user_agent=(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/121.0.0.0 Safari/537.36"
                ),
                viewport={"width": 1280, "height": 800},
                locale="en-US",
                timezone_id="Africa/Lagos",
                extra_http_headers={
                    "Accept-Language": "en-US,en;q=0.9",
                    "Referer": "https://www.google.com/",
                }
            )

            page = await context.new_page()

            await page.goto(
                url,
                wait_until="domcontentloaded",
                timeout=60000
            )

            # Give JS-heavy sites time to settle
            await page.wait_for_timeout(5000)

            screenshot = await page.screenshot(full_page=True)

            await context.close()
            await browser.close()

            return screenshot

    return asyncio.run(run())


if st.button("Capture Screenshot"):
    st.write("Loading page and capturing screenshot...")
    img = get_screenshot(url)

    if img:
        st.image(BytesIO(img), caption="Captured Screenshot", use_column_width=True)
        st.success("Done.")
    else:
        st.error("Screenshot failed.")