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

SAVE_DIR = "/tmp/gradio"  # Ensure files are stored where Gradio can serve them

# Ensure the folder exists
os.makedirs(SAVE_DIR, exist_ok=True)

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

    try:
        wd = webdriver.Chrome(options=options)
        wd.set_window_size(1080, 720)  
        wd.get(url)
        wd.implicitly_wait(20)
        screenshot = wd.get_screenshot_as_png()
    except WebDriverException:
        return "Error: Unable to take screenshot."
    finally:
        if wd:
            wd.quit()

    # Save the screenshot in /tmp/gradio with a unique filename
    filename = f"{SAVE_DIR}/screenshot_{uuid.uuid4().hex}.png"
    image = Image.open(BytesIO(screenshot))
    image.save(filename)

    # Return the file path (Gradio will expose it as a public link)
    return filename

iface = gr.Interface(
    fn=take_screenshot,
    inputs=gr.Textbox(label="Website URL", value="https://kargaranamir.github.io"),
    outputs=gr.File(label="Download Screenshot"),  # Ensures public access
    title="Website Screenshot",
    description="Take a screenshot of a website and download it."
)

iface.launch(share=True)  # Enables external access