import gradio as gr
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from PIL import Image
from io import BytesIO
import time

def take_screenshot(url):
    options = webdriver.ChromeOptions()
    options.add_argument("--headless")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")

    wd = None
    try:
        wd = webdriver.Chrome(options=options)
        wd.set_window_size(1280, 1280)  # Set window size to 1920x1080
        wd.get(url)
        time.sleep(5)
        screenshot = wd.get_screenshot_as_png()

        if len(screenshot) < 1000:
            print("⚠ Screenshot too small, likely error.")
            return None

        image = Image.open(BytesIO(screenshot))

        # Center crop 
        width, height = image.size
        crop_size_side = 1280
        crop_size_top = 1280
        left = (width - crop_size_side) // 2
        top = (height - crop_size_top) // 2
        right = left + 1000
        bottom = top - 500
        cropped_image = image.crop((left, top, right, bottom))

        return image

    except WebDriverException as e:
        print("WebDriverException:", e)
        return None
    finally:
        if wd:
            wd.quit()

iface = gr.Interface(
    fn=take_screenshot,
    inputs=gr.Textbox(label="Website URL", value="https://github.com"),
    outputs=gr.Image(type="pil", label="Cropped Screenshot"),
    title="Website Screenshot",
    description="Takes a 1920x1080 screenshot of the website, center-cropped to 900x900."
)

iface.launch()