app.py

# --- START OF FILE app.py ---

import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from PIL import Image
from io import BytesIO
import tempfile
import time
import os
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
import logging

# ロギング設定 (デバッグに役立つ)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- コアロジック: スクリーンショット生成関数 ---
# (元のコードから変更なし、エラーハンドリングを少し明確化)
def render_fullpage_screenshot(html_code: str, extension_percentage: float) -> Image.Image:
    """
    html_code: HTMLのソースコード文字列
    extension_percentage: 上下に追加する余裕の% (例: 4なら合計4%の余裕)
    戻り値: PIL Imageオブジェクト。エラー時は1x1の黒画像。
    """
    tmp_file = None
    driver = None
    tmp_path = ""

    try:
        # 1) HTMLコードを一時ファイルに保存
        # delete=False と NamedTemporaryFile の close/remove の組み合わせはLinux/Windowsで挙動が安定しないことがあるため、
        # with構文と手動削除を確実に組み合わせる
        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}")

        # 2) ヘッドレスChrome(Chromium)起動オプション
        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") # DPIスケーリング無効化
        options.add_argument("--window-size=1200,800") # 初期ウィンドウサイズ
        # デバッグ用にログレベルを上げる (オプション)
        # options.add_argument("--enable-logging")
        # options.add_argument("--v=1")

        logger.info("Initializing WebDriver...")
        # サービスオブジェクトを作成してログ出力を制御 (オプション)
        # service = webdriver.chrome.service.Service(log_output=subprocess.STDOUT)
        # driver = webdriver.Chrome(options=options, service=service)
        driver = webdriver.Chrome(options=options)
        logger.info("WebDriver initialized.")

        # 3) 初期ウィンドウサイズでページを読み込む (オプションで設定済みなので不要かも)
        # driver.set_window_size(1200, 800)
        file_url = "file://" + tmp_path
        logger.info(f"Navigating to {file_url}")
        driver.get(file_url)

        # 4) ページのロードを待機 (body要素の存在確認)
        logger.info("Waiting for body element...")
        WebDriverWait(driver, 20).until( # タイムアウトを少し延長
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        logger.info("Body element found. Waiting for potential dynamic content loading...")
        time.sleep(3) # 外部リソースやJSによる動的変更の待機時間を少し増やす

        # 5) スクロールバーがキャプチャに写らないようCSSで非表示
        logger.info("Hiding scrollbars via JS...")
        driver.execute_script(
            "document.documentElement.style.overflow = 'hidden';"
            "document.body.style.overflow = 'hidden';"
        )

        # 6) ページ全体の幅・高さを正確に取得
        logger.info("Calculating page dimensions...")
        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"Calculated dimensions: width={scroll_width}, height={scroll_height}")

        # 幅が0や極端に小さい場合はデフォルト値を使うなど対策が必要になる可能性あり
        if scroll_width < 100:
           logger.warning(f"Detected very small scroll_width ({scroll_width}), defaulting to 1200")
           scroll_width = 1200

        # 7) 上下方向にだけユーザー指定の余裕(%)を加えた高さを計算
        adjusted_height = int(scroll_height * (1 + extension_percentage / 100))
        logger.info(f"Adjusted height: {adjusted_height} (extension: {extension_percentage}%)")

        # 8) ウィンドウサイズを、幅はそのまま、縦はadjusted_heightに変更
        logger.info(f"Resizing window to {scroll_width} x {adjusted_height}")
        # set_window_sizeが期待通りに動作しない場合があるため、最大化してからサイズ変更を試すことも検討
        # driver.maximize_window()
        driver.set_window_size(scroll_width, adjusted_height)
        logger.info("Waiting for layout stabilization after resize...")
        time.sleep(3) # レイアウトの安定化待機時間を少し増やす

        # 念のため最上部にスクロール
        logger.info("Scrolling to top...")
        driver.execute_script("window.scrollTo(0, 0)")
        time.sleep(1)

        # 9) スクリーンショット取得
        logger.info("Taking screenshot...")
        png = driver.get_screenshot_as_png()
        logger.info("Screenshot taken successfully.")

        return Image.open(BytesIO(png))

    except Exception as e:
        logger.error(f"Error during screenshot generation: {e}", exc_info=True)
        # エラー時は1x1の黒画像を返す
        return Image.new('RGB', (1, 1), color=(0, 0, 0))
    finally:
        # 必ず WebDriver を終了する
        if driver:
            logger.info("Quitting WebDriver...")
            driver.quit()
            logger.info("WebDriver quit.")
        # 一時ファイルを確実に削除する
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
                logger.info(f"Temporary file removed: {tmp_path}")
            except OSError as e:
                logger.error(f"Error removing temporary file {tmp_path}: {e}")

# --- FastAPI アプリケーション設定 ---
app = FastAPI(
    title="Full Page Screenshot API",
    description="Renders HTML using a headless browser and returns a full-page screenshot.",
    version="1.0.0"
)

# --- APIリクエストモデル ---
class ScreenshotRequest(BaseModel):
    html_code: str = Field(..., description="HTML source code to render.")
    extension_percentage: float = Field(8.0, ge=0, le=100, description="Percentage of extra height to add (top and bottom combined). Default: 8%")

# --- APIエンドポイント ---
@app.post("/screenshot",
          response_class=StreamingResponse,
          tags=["Screenshot"],
          summary="Generate Full Page Screenshot",
          description="Takes HTML code and returns a PNG image of the rendered full page.")
async def create_screenshot(request: ScreenshotRequest):
    """
    Generates a full-page screenshot from the provided HTML code.
    """
    logger.info(f"Received API request. HTML length: {len(request.html_code)}, Extension: {request.extension_percentage}%")
    try:
        pil_image = render_fullpage_screenshot(request.html_code, request.extension_percentage)

        # エラー画像 (1x1) が返ってきたかチェック
        if pil_image.size == (1, 1):
             logger.error("Screenshot generation failed, returning 500 error.")
             # エラーを示すために 500 Internal Server Error を返すことも検討できる
             # raise HTTPException(status_code=500, detail="Failed to generate screenshot. Check HTML or server logs.")
             # ここでは仕様通り 1x1 画像を返すことにするが、APIとしてはエラーの方が親切かもしれない

        # PIL画像をPNGバイトデータに変換
        img_byte_arr = BytesIO()
        pil_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0) # ポインタを先頭に戻す

        logger.info("Returning screenshot as PNG stream.")
        return StreamingResponse(img_byte_arr, media_type="image/png")

    except Exception as e:
        logger.error(f"Unhandled exception in API endpoint: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")

# --- Gradioインターフェース定義 ---
gradio_interface = gr.Interface(
    fn=render_fullpage_screenshot, # コアロジック関数を再利用
    inputs=[
        gr.Textbox(lines=15, label="HTMLコード入力"),
        gr.Slider(minimum=0, maximum=30, step=1.0, value=8, label="上下高さ拡張率（%）") # 最大値を少し増やす
    ],
    outputs=gr.Image(type="pil", label="ページ全体のスクリーンショット"),
    title="Full Page Screenshot (高さ拡張調整可能)",
    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。上下のみユーザー指定の余裕(%)を追加します。APIは /screenshot (POST) で利用可能です。",
    allow_flagging="never" # Hugging Face Spacesでフラグ機能を無効化
)

# --- GradioアプリをFastAPIにマウント ---
# ルートパス ("/") でGradioインターフェースを提供
app = gr.mount_gradio_app(app, gradio_interface, path="/")

# --- uvicorn で実行するための設定 (Hugging Face Spacesでは通常不要) ---
# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=7860)

# --- END OF FILE app.py ---