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

def render_html_to_image(html_code):
    # 入力HTMLを一時ファイルに保存
    tmp_file = tempfile.NamedTemporaryFile(suffix=".html", delete=False)
    tmp_path = tmp_file.name
    tmp_file.write(html_code.encode('utf-8'))
    tmp_file.close()
    
    # Chrome（Chromium）をヘッドレスモードで起動
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')               # ヘッドレス（画面非表示）モード
    options.add_argument('--no-sandbox')             # コンテナ環境でのサンドボックス無効化
    options.add_argument('--disable-dev-shm-usage')  # /dev/shm の共有メモリ領域を不使用
    try:
        driver = webdriver.Chrome(options=options)
        driver.set_window_size(1200, 800)  # ウィンドウサイズを適宜指定
        # ローカルHTMLファイルを開く（file:// スキームを使用）
        driver.get("file://" + tmp_path)
        driver.implicitly_wait(3)         # レンダリング待ち (最大3秒)
        png = driver.get_screenshot_as_png()
    except WebDriverException as e:
        # 失敗時は1x1の透明画像を返すなどの処理
        return Image.new('RGB', (1, 1), color=(0,0,0))
    finally:
        driver.quit()
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
    # PillowでPNGバイナリからImageオブジェクトを生成
    return Image.open(BytesIO(png))

# Gradioインタフェースの定義
iface = gr.Interface(
    fn=render_html_to_image,
    inputs=gr.Textbox(lines=15, label="HTMLコード入力"),
    outputs=gr.Image(type="pil", label="レンダリング結果のスクリーンショット"),
    title="HTMLレンダリング Screenshot アプリ",
    description="入力したHTMLをヘッドレスブラウザでレンダリングしてスクリーンショットを表示します。"
)

if __name__ == "__main__":
    iface.launch()