app.py

# ページのスクリーンショット取得機能を改善した関数
def render_fullpage_screenshot(html_code: str, extension_percentage: float, is_gemini_generated: bool = False) -> Image.Image:
    """
    Renders HTML code to a full-page screenshot using Selenium.

    Args:
        html_code: The HTML source code string.
        extension_percentage: Percentage of extra space to add vertically (e.g., 4 means 4% total).
        is_gemini_generated: Geminiによって生成されたHTMLかどうか（特別な処理を適用）

    Returns:
        A PIL Image object of the screenshot. Returns a 1x1 black image on error.
    """
    tmp_path = None # 初期化
    driver = None # 初期化

    # Gemini生成HTMLの場合は余白を多めに取る
    if is_gemini_generated:
        extension_percentage = max(extension_percentage, 15.0)  # 少なくとも15%の余白を確保

    # 1) Save HTML code to a temporary file
    try:
        with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file:
            tmp_path = tmp_file.name
            tmp_file.write(html_code)
        logger.info(f"HTML saved to temporary file: {tmp_path}")
    except Exception as e:
        logger.error(f"Error writing temporary HTML file: {e}")
        return Image.new('RGB', (1, 1), color=(0, 0, 0)) # エラー時は黒画像

    # 2) Headless Chrome(Chromium) options
    options = Options()
    options.add_argument("--headless")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_argument("--force-device-scale-factor=1")
    # Geminiが生成したHTMLの場合、ウィンドウサイズを大きく設定
    if is_gemini_generated:
        options.add_argument("--window-size=1600,2400")  # 初期サイズを大きく
    
    # Increase logging verbosity for debugging if needed
    # options.add_argument("--enable-logging")
    # options.add_argument("--v=1")

    try:
        logger.info("Initializing WebDriver...")
        driver = webdriver.Chrome(options=options)
        logger.info("WebDriver initialized.")

        # 3) Load page with initial window size
        if is_gemini_generated:
            driver.set_window_size(1600, 2400)  # Gemini生成HTMLには大きなサイズを使用
        else:
            driver.set_window_size(1200, 800)  # 通常のHTML用の標準サイズ
            
        file_url = "file://" + tmp_path
        logger.info(f"Navigating to {file_url}")
        driver.get(file_url)

        # 4) Wait for page load - Gemini生成HTMLは読み込みに時間がかかるので長めに待機
        logger.info("Waiting for body element...")
        wait_timeout = 30 if is_gemini_generated else 15
        WebDriverWait(driver, wait_timeout).until(
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        
        # Font Awesomeとカスタムフォントの読み込み待機（Gemini生成HTMLの場合）
        if is_gemini_generated:
            logger.info("Waiting for potential external resources (Font Awesome, custom fonts)...")
            # Font Awesomeの読み込みを待機（アイコンが表示される）
            time.sleep(5)  # Gemini生成HTMLには長めの待機時間
        else:
            logger.info("Waiting for potential resource loading...")
            time.sleep(3)  # 通常のHTMLには標準の待機時間

        # 5) Hide scrollbars via CSS
        try:
            driver.execute_script(
                "document.documentElement.style.overflow = 'hidden';"
                "document.body.style.overflow = 'hidden';"
            )
            logger.info("Scrollbars hidden via JS.")
        except Exception as e:
            logger.warning(f"Could not hide scrollbars via JS: {e}")

        # 6) Get full page dimensions accurately - Gemini HTMLには追加のJS実行
        try:
            if is_gemini_generated:
                # Gemini生成HTMLの場合、より慎重にコンテンツサイズを計算
                # すべての子要素の合計高さを確認するJSを実行
                scroll_width = driver.execute_script("""
                    return Math.max(
                        document.body.scrollWidth,
                        document.documentElement.scrollWidth,
                        document.body.offsetWidth,
                        document.documentElement.offsetWidth,
                        document.body.clientWidth,
                        document.documentElement.clientWidth
                    );
                """)
                
                scroll_height = driver.execute_script("""
                    // ドキュメント全体の高さを取得
                    var docHeight = Math.max(
                        document.body.scrollHeight,
                        document.documentElement.scrollHeight,
                        document.body.offsetHeight,
                        document.documentElement.offsetHeight,
                        document.body.clientHeight,
                        document.documentElement.clientHeight
                    );
                    
                    // すべての要素をチェックして最大のY位置を見つける
                    var allElements = document.getElementsByTagName('*');
                    var maxBottom = 0;
                    
                    for (var i = 0; i < allElements.length; i++) {
                        var rect = allElements[i].getBoundingClientRect();
                        var bottom = rect.bottom + window.scrollY;
                        if (bottom > maxBottom) {
                            maxBottom = bottom;
                        }
                    }
                    
                    // 最大のY位置とドキュメント高さの大きい方を返す
                    return Math.max(docHeight, maxBottom);
                """)
            else:
                # 通常のHTML用の標準の計算
                scroll_width = driver.execute_script(
                    "return Math.max(document.body.scrollWidth, document.documentElement.scrollWidth, document.body.offsetWidth, document.documentElement.offsetWidth)"
                )
                scroll_height = driver.execute_script(
                    "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight, document.body.offsetHeight, document.documentElement.offsetHeight)"
                )
                
            logger.info(f"Detected dimensions: width={scroll_width}, height={scroll_height}")
            # Ensure minimum dimensions to avoid errors
            scroll_width = max(scroll_width, 100) # 最小幅を設定
            scroll_height = max(scroll_height, 100) # 最小高さを設定

        except Exception as e:
            logger.error(f"Error getting page dimensions: {e}")
            # フォールバックとしてデフォルト値を設定
            if is_gemini_generated:
                scroll_width = 1600
                scroll_height = 2400
            else:
                scroll_width = 1200
                scroll_height = 800
            logger.warning(f"Falling back to dimensions: width={scroll_width}, height={scroll_height}")

        # 7) Calculate adjusted height with user-specified margin
        adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
        # Ensure adjusted height is not excessively large or small
        adjusted_height = max(adjusted_height, scroll_height, 100) # 最小高さを確保
        logger.info(f"Adjusted height calculated: {adjusted_height} (extension: {extension_percentage}%)")

        # 8) Set window size to full page dimensions (width) and adjusted height
        logger.info(f"Resizing window to: width={scroll_width}, height={adjusted_height}")
        driver.set_window_size(scroll_width, adjusted_height)
        logger.info("Waiting for layout stabilization after resize...")
        wait_time = 5 if is_gemini_generated else 3  # Gemini生成HTMLには長めの待機時間
        time.sleep(wait_time)

        # Scroll to top just in case
        try:
            driver.execute_script("window.scrollTo(0, 0)")
            time.sleep(1)
            logger.info("Scrolled to top.")
        except Exception as e:
             logger.warning(f"Could not scroll to top: {e}")

        # 9) Take screenshot
        logger.info("Taking screenshot...")
        # Gemini生成HTMLの場合、完全なスクリーンショットを取得するための特別な処理
        if is_gemini_generated:
            # スクロールしながら複数の部分スクリーンショットを撮り、それらを結合する方法もあります
            # ここではシンプルに全体のスクリーンショットを取得します
            time.sleep(1)  # 最終的な安定化のための短い待機
            
        png = driver.get_screenshot_as_png()
        logger.info("Screenshot taken successfully.")

        # Convert to PIL Image
        img = Image.open(BytesIO(png))
        return img

    except Exception as e:
        logger.error(f"An error occurred during screenshot generation: {e}", exc_info=True)
        return Image.new('RGB', (1, 1), color=(0, 0, 0)) # Return black 1x1 image on error
    finally:
        logger.info("Cleaning up...")
        if driver:
            try:
                driver.quit()
                logger.info("WebDriver quit successfully.")
            except Exception as e:
                logger.error(f"Error quitting WebDriver: {e}")
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
                logger.info(f"Temporary file {tmp_path} removed.")
            except Exception as e:
                logger.error(f"Error removing temporary file {tmp_path}: {e}")

# --- Geminiを使った新しい関数 ---
def text_to_screenshot(text: str, extension_percentage: float) -> Image.Image:
    """テキストをGemini APIでHTMLに変換し、スクリーンショットを生成する統合関数"""
    try:
        # 1. テキストからHTMLを生成
        html_code = generate_html_from_text(text)
        
        # 2. HTMLからスクリーンショットを生成（Gemini生成フラグをTrue）
        return render_fullpage_screenshot(html_code, extension_percentage, is_gemini_generated=True)
    except Exception as e:
        logger.error(f"テキストからスクリーンショット生成中にエラーが発生: {e}", exc_info=True)
        return Image.new('RGB', (1, 1), color=(0, 0, 0))  # エラー時は黒画像