import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import WebDriverException, TimeoutException
import tempfile
import os
from PIL import Image
import time
import base64
from bs4 import BeautifulSoup  # HTML解析用


def take_screenshot_from_html(html_code, height_adjustment_percent=3):
    options = Options()
    options.add_argument("--headless=new")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_argument("--disable-gpu")
    options.add_argument("--window-position=-2400,-2400")

    try:
        driver = webdriver.Chrome(options=options)

        # BeautifulSoupでHTMLを解析
        soup = BeautifulSoup(html_code, 'html.parser')

        # <head>タグがあるか確認し、なければ追加
        if not soup.head:
            head = soup.new_tag("head")
            soup.insert(0, head)

        # <meta charset>タグがあるか確認し、なければ追加（UTF-8を指定）
        if not soup.head.find("meta", charset=True):
            meta_charset = soup.new_tag("meta")
            meta_charset["charset"] = "UTF-8"
            soup.head.insert(0, meta_charset)

        # 修正したHTMLを文字列に戻す
        modified_html = str(soup)

        html_base64 = base64.b64encode(modified_html.encode('utf-8')).decode('utf-8')
        driver.get(f"data:text/html;base64,{html_base64}")

        time.sleep(3)

        width = driver.execute_script(
            "return Math.max(document.body.scrollWidth, document.body.offsetWidth, "
            "document.documentElement.clientWidth, document.documentElement.scrollWidth, "
            "document.documentElement.offsetWidth);"
        )
        original_height = driver.execute_script(
            "return Math.max(document.body.scrollHeight, document.body.offsetHeight, "
            "document.documentElement.clientHeight, document.documentElement.scrollHeight, "
            "document.documentElement.offsetHeight);"
        )

        adjusted_height = int(original_height * (1 + height_adjustment_percent / 100))
        driver.set_window_size(min(width + 100, 1920), min(adjusted_height + 100, 5000))
        time.sleep(1)

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp:
            temp_filename = temp.name
        driver.save_screenshot(temp_filename)

        driver.quit()
        image = Image.open(temp_filename)
        os.unlink(temp_filename)
        return image, "スクリーンショットを取得しました。"

    except (WebDriverException, TimeoutException) as e:
        driver.quit()
        return None, f"エラー: {str(e)}"
    except Exception as e:
        driver.quit()
        return None, f"予期しないエラー: {str(e)}"


with gr.Blocks(title="HTML Screenshot Tool") as demo:
    gr.Markdown("# HTML スクリーンショットツール")
    gr.Markdown("HTMLコードを入力して、レンダリングされたページのフルスクリーンショットを取得します。")

    html_input = gr.Textbox(label="HTMLコード", placeholder="<p>Hello, world!</p>", lines=5)
    height_adjustment_slider = gr.Slider(minimum=0, maximum=10, value=3, step=0.1, label="高さ調整率 (%)")
    html_button = gr.Button("スクリーンショットを取得")
    html_image_out = gr.Image(label="スクリーンショット")
    html_text_out = gr.Textbox(label="結果")

    html_button.click(
        take_screenshot_from_html,
        inputs=[html_input, height_adjustment_slider],
        outputs=[html_image_out, html_text_out]
    )

demo.launch()