import gradio as gr
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from PIL import Image
from io import BytesIO

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

    wd = None
    try:
        service = Service(ChromeDriverManager().install())
        wd = webdriver.Chrome(service=service, options=options)
        wd.get(url)
        wd.implicitly_wait(10)

        if capture_type == "전체 페이지":
            # 전체 페이지 높이와 너비를 가져옵니다.
            total_width = wd.execute_script("return document.body.offsetWidth")
            total_height = wd.execute_script("return document.body.parentNode.scrollHeight")
            wd.set_window_size(total_width, total_height)  # 윈도우 크기를 전체 페이지 크기로 조정합니다.
        else:
            wd.set_window_size(1024, 768)  # 기본 크기

        screenshot = wd.get_screenshot_as_png()
    except WebDriverException as e:
        return Image.new('RGB', (1, 1))
    finally:
        if wd:
            wd.quit()

    return Image.open(BytesIO(screenshot))

iface = gr.Interface(
    fn=take_screenshot,
    inputs=[
        gr.inputs.Textbox(label="Website URL", default="https://korating.com"),
        gr.inputs.Radio(choices=["기본 크기", "전체 페이지"], label="캡처 유형 선택")
    ],
    outputs=gr.Image(type="pil", tool="editor"),
    title="Website Screenshot",
    description="Take a screenshot of a website. Choose '전체 페이지' to capture the full page.",
)

iface.launch()