from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
import gradio as gr

def capture_webpage_screenshot(url):
    save_path = "screenshot.png"
    # ChromeDriver를 자동으로 설치하고 초기화합니다.
    service = Service(ChromeDriverManager().install())
    options = webdriver.ChromeOptions()
    # 필요한 경우, headless 모드로 설정할 수 있습니다.
    options.add_argument('--headless')
    options.add_argument('--disable-gpu')  # GPU 사용 비활성화
    options.add_argument('window-size=1200x600')  # 윈도우 크기 설정
    
    # WebDriver 초기화
    driver = webdriver.Chrome(service=service, options=options)
    
    # 주어진 URL로 이동
    driver.get(url)
    
    # 스크린샷 캡처 및 저장
    driver.save_screenshot(save_path)
    
    # WebDriver 종료
    driver.quit()
    return save_path


with gr.Blocks() as demo:
    txt = gr.Textbox()

    btn = gr.Button()
    img = gr.Image()
    btn.click(capture_webpage_screenshot, [txt], [img])


demo.launch()