import gradio as gr
from selenium import webdriver
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.chrome.service import Service
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 = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)

    if capture_type == "전체 페이지":
        # 전체 페이지 캡처 로직 (예시: JavaScript를 사용하여 페이지 높이 측정)
        wd.get(url)
        total_height = wd.execute_script("return document.body.parentNode.scrollHeight")
        wd.set_window_size(1024, total_height)  # 창 크기를 전체 페이지 높이에 맞춤
    else:
        wd.set_window_size(1024, 768)  # 기본 크기

    wd.get(url)
    screenshot = wd.get_screenshot_as_png()
    wd.quit()

    return Image.open(BytesIO(screenshot))

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

iface.launch()