app.py

import gradio as gr
from fastapi import FastAPI, HTTPException, Body
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
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
import logging
import numpy as np
import threading
import queue
import uuid
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from huggingface_hub import hf_hub_download, upload_file, login

# 正しいGemini関連のインポート
import google.generativeai as genai

# ロギング設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- HuggingFace Hub アップロード機能 ---
class HuggingFaceUploader:
    """HuggingFace Hubへ画像をアップロードする機能を提供するクラス"""
    def __init__(self):
        self.repo_id = os.environ.get("HF_REPO_ID", "leave-everything/ChoTensaiJinGrareko")
        self.token = os.environ.get("HF_TOKEN", None)
        if self.token:
            try:
                login(token=self.token)
                logger.info(f"HuggingFace Hubにログインしました。リポジトリ: {self.repo_id}")
            except Exception as e:
                logger.error(f"HuggingFace Hubへのログインに失敗: {e}")
        else:
            logger.warning("HF_TOKEN環境変数が設定されていません。アップロード機能は制限されます。")
    
    def upload_image(self, image, prefix="generated"):
        """
        PIL Imageをアップロードし、アクセス可能なURLを返す
        
        Args:
            image: PIL.Image - アップロードする画像
            prefix: str - ファイル名のプレフィックス
            
        Returns:
            str - アップロードされた画像のURL
        """
        try:
            if not self.token:
                logger.error("HF_TOKENが設定されていないため、アップロードできません")
                return None
                
            # ユニークなファイル名を生成 - JPEGフォーマットを使用
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            unique_id = str(uuid.uuid4())[:8]
            filename = f"{prefix}_{timestamp}_{unique_id}.jpg"
            path_in_repo = f"images/{filename}"
            
            # 一時ファイルに保存 - JPEGフォーマットで保存
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
                tmp_path = tmp_file.name
                image.save(tmp_path, format="JPEG", quality=95)  # 高品質JPEGとして保存
            
            logger.info(f"画像を一時ファイルに保存: {tmp_path}")
            
            # HuggingFaceにアップロード
            logger.info(f"HuggingFace Hubにアップロード中: {path_in_repo}")
            upload_info = upload_file(
                path_or_fileobj=tmp_path,
                path_in_repo=path_in_repo,
                repo_id=self.repo_id,
                repo_type="dataset"
            )
            
            # 一時ファイルを削除
            try:
                os.remove(tmp_path)
            except Exception as e:
                logger.warning(f"一時ファイル削除エラー: {e}")
            
            # URLを構築して返す - データセットリポジトリの場合は datasets/ を含める
            url = f"https://huggingface.co/datasets/{self.repo_id}/resolve/main/{path_in_repo}"
            logger.info(f"アップロード成功: {url}")
            return url
            
        except Exception as e:
            logger.error(f"HuggingFace Hubへのアップロード中にエラー: {e}", exc_info=True)
            return None

# グローバルなアップローダーインスタンスを作成
hf_uploader = HuggingFaceUploader()

# --- WebDriverプールの実装 ---
class WebDriverPool:
    """WebDriverインスタンスを再利用するためのプール"""
    def __init__(self, max_drivers=3):
        self.driver_queue = queue.Queue()
        self.max_drivers = max_drivers
        self.lock = threading.Lock()
        self.count = 0
        logger.info(f"WebDriverプールを初期化: 最大 {max_drivers} ドライバー")
    
    def get_driver(self):
        """プールからWebDriverを取得、なければ新規作成"""
        if not self.driver_queue.empty():
            logger.info("既存のWebDriverをプールから取得")
            return self.driver_queue.get()
        
        with self.lock:
            if self.count < self.max_drivers:
                self.count += 1
                logger.info(f"新しいWebDriverを作成 (合計: {self.count}/{self.max_drivers})")
                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")
                options.add_argument("--disable-features=NetworkService")
                options.add_argument("--dns-prefetch-disable")
                
                # 環境変数からWebDriverパスを取得（任意）
                webdriver_path = os.environ.get("CHROMEDRIVER_PATH")
                if webdriver_path and os.path.exists(webdriver_path):
                    logger.info(f"CHROMEDRIVER_PATH使用: {webdriver_path}")
                    service = webdriver.ChromeService(executable_path=webdriver_path)
                    return webdriver.Chrome(service=service, options=options)
                else:
                    logger.info("デフォルトのChromeDriverを使用")
                    return webdriver.Chrome(options=options)
        
        # 最大数に達した場合は待機
        logger.info("WebDriverプールがいっぱいです。利用可能なドライバーを待機中...")
        return self.driver_queue.get()
    
    def release_driver(self, driver):
        """ドライバーをプールに戻す"""
        if driver:
            try:
                # ブラウザをリセット
                driver.get("about:blank")
                driver.execute_script("""
                    document.documentElement.style.overflow = '';
                    document.body.style.overflow = '';
                """)
                self.driver_queue.put(driver)
                logger.info("WebDriverをプールに戻しました")
            except Exception as e:
                logger.error(f"ドライバーをプールに戻す際にエラー: {e}")
                driver.quit()
                with self.lock:
                    self.count -= 1
    
    def close_all(self):
        """全てのドライバーを終了"""
        logger.info("WebDriverプールを終了します")
        closed = 0
        while not self.driver_queue.empty():
            try:
                driver = self.driver_queue.get(block=False)
                driver.quit()
                closed += 1
            except queue.Empty:
                break
            except Exception as e:
                logger.error(f"ドライバー終了中にエラー: {e}")
        
        logger.info(f"{closed}個のWebDriverを終了しました")
        with self.lock:
            self.count = 0

# グローバルなWebDriverプールを作成
# サーバー環境のリソースに合わせて調整
driver_pool = WebDriverPool(max_drivers=int(os.environ.get("MAX_WEBDRIVERS", "3")))

# --- Gemini統合 ---
class GeminiRequest(BaseModel):
    """Geminiへのリクエストデータモデル"""
    text: str
    extension_percentage: float = 15.0  # デフォルト値15%
    temperature: float = 1.0  # デフォルト値を1.0に設定
    trim_whitespace: bool = True  # 余白トリミングオプション（デフォルト有効）
    style: str = "standard"  # デフォルトはstandard

class ScreenshotRequest(BaseModel):
    """スクリーンショットリクエストモデル"""
    html_code: str
    extension_percentage: float = 15.0  # デフォルト値15%
    trim_whitespace: bool = True  # 余白トリミングオプション（デフォルト有効）
    style: str = "standard"  # デフォルトはstandard

# --- レスポンスモデル ---
class ImageUrlResponse(BaseModel):
    """画像URLのレスポンスモデル"""
    url: str
    status: str = "success"
    message: str = "画像が正常に生成されました"

# HTMLのFont Awesomeレイアウトを改善する関数 - プリロード機能を追加
def enhance_font_awesome_layout(html_code):
    """Font Awesomeレイアウトを改善し、プリロードタグを追加"""
    # Font Awesomeリソースのプリロード - パフォーマンス向上
    fa_preload = """
    <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-solid-900.woff2" as="font" type="font/woff2" crossorigin>
    <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-regular-400.woff2" as="font" type="font/woff2" crossorigin>
    <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-brands-400.woff2" as="font" type="font/woff2" crossorigin>
    """
    
    # CSSを追加
    fa_fix_css = """
    <style>
    /* Font Awesomeアイコンのレイアウト修正 */
    [class*="fa-"] {
        display: inline-block !important;
        margin-right: 8px !important;
        vertical-align: middle !important;
    }
    
    /* テキスト内のアイコン位置調整 */
    h1 [class*="fa-"], h2 [class*="fa-"], h3 [class*="fa-"], 
    h4 [class*="fa-"], h5 [class*="fa-"], h6 [class*="fa-"] {
        vertical-align: middle !important;
        margin-right: 10px !important;
    }
    
    /* 特定パターンの修正 */
    .fa + span, .fas + span, .far + span, .fab + span,
    span + .fa, span + .fas, span + .far + span {
        display: inline-block !important;
        margin-left: 5px !important;
    }
    
    /* カード内アイコン修正 */
    .card [class*="fa-"], .card-body [class*="fa-"] {
        float: none !important;
        clear: none !important;
        position: relative !important;
    }
    
    /* アイコンと文字が重なる場合の調整 */
    li [class*="fa-"], p [class*="fa-"] {
        margin-right: 10px !important;
    }
    
    /* インラインアイコンのスペーシング */
    .inline-icon {
        display: inline-flex !important;
        align-items: center !important;
        justify-content: flex-start !important;
    }
    
    /* アイコン後のテキスト */
    [class*="fa-"] + span {
        display: inline-block !important;
        vertical-align: middle !important;
    }
    </style>
    """

    # headタグがある場合はその中に追加
    if '<head>' in html_code:
        return html_code.replace('</head>', f'{fa_preload}{fa_fix_css}</head>')
    # HTMLタグがある場合はその後に追加
    elif '<html' in html_code:
        head_end = html_code.find('</head>')
        if head_end > 0:
            return html_code[:head_end] + fa_preload + fa_fix_css + html_code[head_end:]
        else:
            body_start = html_code.find('<body')
            if body_start > 0:
                return html_code[:body_start] + f'<head>{fa_preload}{fa_fix_css}</head>' + html_code[body_start:]

    # どちらもない場合は先頭に追加
    return f'<html><head>{fa_preload}{fa_fix_css}</head>' + html_code + '</html>'

def load_system_instruction(style="standard"):
    """
    指定されたスタイルのシステムインストラクションを読み込む

    Args:
        style: 使用するスタイル名 (standard, cute, resort, cool, dental, school)

    Returns:
        読み込まれたシステムインストラクション
    """
    try:
        # 有効なスタイル一覧
        valid_styles = ["standard", "cute", "resort", "cool", "dental", "school","KOKUGO"]

        # スタイルの検証
        if style not in valid_styles:
            logger.warning(f"無効なスタイル '{style}' が指定されました。デフォルトの 'standard' を使用します。")
            style = "standard"

        logger.info(f"スタイル '{style}' のシステムインストラクションを読み込みます")

        # まず、ローカルのスタイルディレクトリ内のprompt.txtを確認
        local_path = os.path.join(os.path.dirname(__file__), style, "prompt.txt")

        # ローカルファイルが存在する場合はそれを使用
        if os.path.exists(local_path):
            logger.info(f"ローカルファイルを使用: {local_path}")
            with open(local_path, 'r', encoding='utf-8') as file:
                instruction = file.read()
            return instruction

        # HuggingFaceリポジトリからのファイル読み込みを試行
        try:
            # スタイル固有のファイルパスを指定
            file_path = hf_hub_download(
                repo_id="tomo2chin2/GURAREKOstlyle",
                filename=f"{style}/prompt.txt",
                repo_type="dataset"
            )

            logger.info(f"スタイル '{style}' のプロンプトをHuggingFaceから読み込みました: {file_path}")
            with open(file_path, 'r', encoding='utf-8') as file:
                instruction = file.read()
            return instruction

        except Exception as style_error:
            # スタイル固有ファイルの読み込みに失敗した場合、デフォルトのprompt.txtを使用
            logger.warning(f"スタイル '{style}' のプロンプト読み込みに失敗: {str(style_error)}")
            logger.info("デフォルトのprompt.txtを読み込みます")

            file_path = hf_hub_download(
                repo_id="leave-everything/GURAREKOstyle",
                filename="prompt.txt",
                repo_type="dataset"
            )

            with open(file_path, 'r', encoding='utf-8') as file:
                instruction = file.read()

            logger.info("デフォルトのシステムインストラクションを読み込みました")
            return instruction

    except Exception as e:
        error_msg = f"システムインストラクションの読み込みに失敗: {str(e)}"
        logger.error(error_msg)
        raise ValueError(error_msg)

def generate_html_from_text(text, temperature=0.3, style="standard"):
    """テキストからHTMLを生成する"""
    try:
        # APIキーの取得と設定
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            logger.error("GEMINI_API_KEY 環境変数が設定されていません")
            raise ValueError("GEMINI_API_KEY 環境変数が設定されていません")

        # モデル名の取得（環境変数から、なければデフォルト値）
        model_name = os.environ.get("GEMINI_MODEL", "gemini-flash-latest")
        logger.info(f"使用するGeminiモデル: {model_name}")

        # Gemini APIの設定
        genai.configure(api_key=api_key)

        # 指定されたスタイルのシステムインストラクションを読み込む
        system_instruction = load_system_instruction(style)

        # モデル初期化
        logger.info(f"Gemini APIにリクエストを送信: テキスト長さ = {len(text)}, 温度 = {temperature}, スタイル = {style}")

        # モデル初期化
        model = genai.GenerativeModel(model_name)

        # 生成設定 - ばらつきを減らすために設定を調整
        generation_config = {
            "temperature": temperature,  # より低い温度を設定
            "top_p": 0.7,                # 0.95から0.7に下げて出力の多様性を制限
            "top_k": 20,                 # 64から20に下げて候補を絞る
            "max_output_tokens": 8192,
            "candidate_count": 1         # 候補は1つだけ生成
        }

        # 安全設定 - デフォルトの安全設定を使用
        safety_settings = [
            {
                "category": "HARM_CATEGORY_HARASSMENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
            },
            {
                "category": "HARM_CATEGORY_HATE_SPEECH",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
            },
            {
                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
            },
            {
                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
            }
        ]

        # プロンプト構築
        prompt = f"{system_instruction}\n\n{text}"

        # コンテンツ生成
        response = model.generate_content(
            prompt,
            generation_config=generation_config,
            safety_settings=safety_settings
        )

        # レスポンスからHTMLを抽出
        raw_response = response.text

        # HTMLタグ部分だけを抽出（```html と ``` の間）
        html_start = raw_response.find("```html")
        html_end = raw_response.rfind("```")

        if html_start != -1 and html_end != -1 and html_start < html_end:
            html_start += 7  # "```html" の長さ分進める
            html_code = raw_response[html_start:html_end].strip()
            logger.info(f"HTMLの生成に成功: 長さ = {len(html_code)}")

            # Font Awesomeのレイアウト改善
            html_code = enhance_font_awesome_layout(html_code)
            logger.info("Font Awesomeレイアウトの最適化を適用しました")

            return html_code
        else:
            # HTMLタグが見つからない場合、レスポンス全体を返す
            logger.warning("レスポンスから ```html ``` タグが見つかりませんでした。全テキストを返します。")
            return raw_response

    except Exception as e:
        logger.error(f"HTML生成中にエラーが発生: {e}", exc_info=True)
        raise Exception(f"Gemini APIでのHTML生成に失敗しました: {e}")

# 画像から余分な空白領域をトリミングする関数 - NumPyを使って最適化
def trim_image_whitespace(image, threshold=250, padding=10):
    """
    NumPyを使用して最適化された画像トリミング関数
    
    Args:
        image: PIL.Image - 入力画像
        threshold: int - どの明るさ以上を空白と判断するか (0-255)
        padding: int - トリミング後に残す余白のピクセル数

    Returns:
        トリミングされたPIL.Image
    """
    try:
        # グレースケールに変換
        gray = image.convert('L')
        
        # NumPy配列として取得（高速処理のため）
        np_image = np.array(gray)
        
        # マスク作成（非白ピクセル）
        mask = np_image < threshold
        
        # マスクから行と列のインデックスを取得
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)
        
        # 非空のインデックス範囲を取得
        if np.any(rows) and np.any(cols):
            row_indices = np.where(rows)[0]
            col_indices = np.where(cols)[0]
            
            # 範囲取得
            min_y, max_y = row_indices[0], row_indices[-1]
            min_x, max_x = col_indices[0], col_indices[-1]
            
            # パディング追加
            min_x = max(0, min_x - padding)
            min_y = max(0, min_y - padding)
            max_x = min(image.width - 1, max_x + padding)
            max_y = min(image.height - 1, max_y + padding)
            
            # 画像をトリミング
            trimmed = image.crop((min_x, min_y, max_x + 1, max_y + 1))
            
            logger.info(f"画像をトリミングしました: 元サイズ {image.width}x{image.height} → トリミング後 {trimmed.width}x{trimmed.height}")
            return trimmed
        
        logger.warning("トリミング領域が見つかりません。元の画像を返します。")
        return image
        
    except Exception as e:
        logger.error(f"画像トリミング中にエラー: {e}", exc_info=True)
        return image  # エラー時は元の画像を返す

# 最適化されたスクリーンショット生成関数 - 外部から初期化済みドライバーを受け取れるように
def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0,
                              trim_whitespace: bool = True, driver=None) -> Image.Image:
    """
    Renders HTML code to a full-page screenshot using Selenium.
    Optimized to accept an external driver or get one from the pool.

    Args:
        html_code: The HTML source code string.
        extension_percentage: Percentage of extra space to add vertically.
        trim_whitespace: Whether to trim excess whitespace from the image.
        driver: An optional pre-initialized WebDriver instance.

    Returns:
        A PIL Image object of the screenshot.
    """
    tmp_path = None
    driver_from_pool = False
    
    # ドライバーがない場合はプールから取得
    if driver is None:
        driver = driver_pool.get_driver()
        driver_from_pool = True
        logger.info("WebDriverプールからドライバーを取得しました")

    # 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}")
        if driver_from_pool:
            driver_pool.release_driver(driver)
        return Image.new('RGB', (1, 1), color=(0, 0, 0))

    try:
        # ウィンドウサイズ初期設定
        initial_width = 1200
        initial_height = 1000
        driver.set_window_size(initial_width, initial_height)
        file_url = "file://" + tmp_path
        logger.info(f"Navigating to {file_url}")
        driver.get(file_url)

        # ページ読み込み待機 - 動的な待機時間を実装
        logger.info("Waiting for body element...")
        WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        logger.info("Body element found. Waiting for resource loading...")

        # リソース読み込みの動的待機 - 最適化
        max_wait = 5  # 最大待機時間（秒）
        wait_increment = 0.2  # 確認間隔
        wait_time = 0
        
        while wait_time < max_wait:
            resource_state = driver.execute_script("""
                return {
                    complete: document.readyState === 'complete',
                    imgCount: document.images.length,
                    imgLoaded: Array.from(document.images).filter(img => img.complete).length,
                    faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length
                };
            """)
            
            # ドキュメント完了かつ画像が読み込まれている場合、待機終了
            if resource_state['complete'] and (resource_state['imgCount'] == 0 or 
                                             resource_state['imgLoaded'] == resource_state['imgCount']):
                logger.info(f"リソース読み込み完了: {resource_state}")
                break
                
            time.sleep(wait_increment)
            wait_time += wait_increment
            logger.info(f"リソース待機中... {wait_time:.1f}秒経過, 状態: {resource_state}")
        
        # Font Awesome要素が多い場合は追加待機
        fa_count = resource_state.get('faElements', 0)
        if fa_count > 30:
            logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...")
            time.sleep(min(1.0, fa_count / 100))  # 要素数に応じて待機（最大1秒）

        # コンテンツレンダリングのためのスクロール処理 - パフォーマンス改善
        logger.info("Performing content rendering scroll...")
        total_height = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight);")
        viewport_height = driver.execute_script("return window.innerHeight;")
        scrolls_needed = max(1, min(5, total_height // viewport_height))  # 最大5回までに制限
        
        # スクロール処理の高速化
        for i in range(scrolls_needed):
            scroll_pos = i * (viewport_height - 100)  # 少しだけオーバーラップ
            driver.execute_script(f"window.scrollTo(0, {scroll_pos});")
            time.sleep(0.1)  # 高速化のため待機時間短縮
        
        # トップに戻る
        driver.execute_script("window.scrollTo(0, 0);")
        time.sleep(0.2)  # 短い待機に変更
        logger.info("Scroll rendering completed")

        # スクロールバーを非表示に
        driver.execute_script("""
            document.documentElement.style.overflow = 'hidden';
            document.body.style.overflow = 'hidden';
        """)
        
        # ページの寸法を取得
        dimensions = driver.execute_script("""
            return {
                width: Math.max(
                    document.documentElement.scrollWidth,
                    document.documentElement.offsetWidth,
                    document.documentElement.clientWidth,
                    document.body ? document.body.scrollWidth : 0,
                    document.body ? document.body.offsetWidth : 0,
                    document.body ? document.body.clientWidth : 0
                ),
                height: Math.max(
                    document.documentElement.scrollHeight,
                    document.documentElement.offsetHeight,
                    document.documentElement.clientHeight,
                    document.body ? document.body.scrollHeight : 0,
                    document.body ? document.body.offsetHeight : 0,
                    document.body ? document.body.clientHeight : 0
                )
            };
        """)
        scroll_width = dimensions['width']
        scroll_height = dimensions['height']
        logger.info(f"Detected dimensions: width={scroll_width}, height={scroll_height}")

        # 最小/最大値の設定
        scroll_width = max(scroll_width, 100)
        scroll_height = max(scroll_height, 100)
        scroll_width = min(scroll_width, 2000)
        scroll_height = min(scroll_height, 4000)
        
        # レイアウト安定化のための待機 
        time.sleep(2.0)  

        # 高さに余白を追加
        adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
        adjusted_height = max(adjusted_height, scroll_height, 100)
        
        # ウィンドウサイズを調整
        adjusted_width = scroll_width
        logger.info(f"Resizing window to: width={adjusted_width}, height={adjusted_height}")
        driver.set_window_size(adjusted_width, adjusted_height)
        time.sleep(0.5)  # 短縮した待機時間

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

        # PIL画像に変換
        img = Image.open(BytesIO(png))
        logger.info(f"Screenshot dimensions: {img.width}x{img.height}")

        # 余白トリミング - 最適化版を使用
        if trim_whitespace:
            img = trim_image_whitespace(img, threshold=248, padding=20)
            logger.info(f"Trimmed dimensions: {img.width}x{img.height}")

        return img

    except Exception as e:
        logger.error(f"Error during screenshot generation: {e}", exc_info=True)
        # エラー時は小さい黒画像を返す
        return Image.new('RGB', (1, 1), color=(0, 0, 0))
    finally:
        logger.info("Cleaning up...")
        # WebDriverプールに戻す
        if driver_from_pool:
            driver_pool.release_driver(driver)
            logger.info("Returned driver to pool")
        # 一時ファイル削除
        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}")

# --- 並列処理を活用した新しい関数 ---
def text_to_screenshot_parallel(text: str, extension_percentage: float, temperature: float = 0.3,
                             trim_whitespace: bool = True, style: str = "standard") -> tuple:
    """
    テキストをGemini APIでHTMLに変換し、並列処理でスクリーンショットを生成する関数
    
    Returns:
        tuple - (PIL.Image, URL) - 生成された画像とHuggingFaceのURL
    """
    start_time = time.time()
    logger.info("並列処理によるテキスト→スクリーンショット生成を開始")
    
    try:
        # WebDriverと HTML生成を並列で実行
        with ThreadPoolExecutor(max_workers=2) as executor:
            # Gemini APIリクエストタスク
            html_future = executor.submit(
                generate_html_from_text, 
                text=text, 
                temperature=temperature, 
                style=style
            )
            
            # WebDriver初期化タスク - プール使用
            driver_future = executor.submit(driver_pool.get_driver)
            
            # 結果を取得
            html_code = html_future.result()
            driver = driver_future.result()
            
        # ドライバーはプールから取得しているためフラグ設定
        driver_from_pool = True
        
        # HTMLコードとドライバーが準備できたらスクリーンショット生成
        logger.info(f"HTML生成完了：{len(html_code)}文字。スクリーンショット生成開始。")
        
        # レンダリング前にドライバーの初期設定
        tmp_path = None
        try:
            # 一時ファイルにHTMLを保存
            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を一時ファイルに保存: {tmp_path}")
            
            # ウィンドウサイズ初期設定
            initial_width = 1200
            initial_height = 1000
            driver.set_window_size(initial_width, initial_height)
            file_url = "file://" + tmp_path
            logger.info(f"ページに移動: {file_url}")
            driver.get(file_url)
            
            # ここからスクリーンショット生成ロジック（前の実装と同様）
            # ページ読み込み待機 - 動的な待機時間を実装
            logger.info("body要素を待機...")
            WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )
            logger.info("body要素を検出。リソース読み込みを待機...")
            
            # リソース読み込みの動的待機 - 最適化
            max_wait = 3  # 最大待機時間（秒）
            wait_increment = 0.2  # 確認間隔
            wait_time = 0
            
            while wait_time < max_wait:
                resource_state = driver.execute_script("""
                    return {
                        complete: document.readyState === 'complete',
                        imgCount: document.images.length,
                        imgLoaded: Array.from(document.images).filter(img => img.complete).length,
                        faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length
                    };
                """)
                
                # ドキュメント完了かつ画像が読み込まれている場合、待機終了
                if resource_state['complete'] and (resource_state['imgCount'] == 0 or 
                                                 resource_state['imgLoaded'] == resource_state['imgCount']):
                    logger.info(f"リソース読み込み完了: {resource_state}")
                    break
                    
                time.sleep(wait_increment)
                wait_time += wait_increment
            
            # Font Awesome要素が多い場合は追加待機
            fa_count = resource_state.get('faElements', 0)
            if fa_count > 30:
                logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...")
                time.sleep(min(1.0, fa_count / 100))  # 要素数に応じて待機（最大0.5秒）
            
            # コンテンツレンダリングのための簡易スクロール
            driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(0.2)
            driver.execute_script("window.scrollTo(0, 0);")
            time.sleep(0.2)
            
            # スクロールバーを非表示に
            driver.execute_script("""
                document.documentElement.style.overflow = 'hidden';
                document.body.style.overflow = 'hidden';
            """)
            
            # ページの寸法を取得
            dimensions = driver.execute_script("""
                return {
                    width: Math.max(
                        document.documentElement.scrollWidth,
                        document.documentElement.offsetWidth,
                        document.documentElement.clientWidth,
                        document.body ? document.body.scrollWidth : 0,
                        document.body ? document.body.offsetWidth : 0,
                        document.body ? document.body.clientWidth : 0
                    ),
                    height: Math.max(
                        document.documentElement.scrollHeight,
                        document.documentElement.offsetHeight,
                        document.documentElement.clientHeight,
                        document.body ? document.body.scrollHeight : 0,
                        document.body ? document.body.offsetHeight : 0,
                        document.body ? document.body.clientHeight : 0
                    )
                };
            """)
            scroll_width = dimensions['width']
            scroll_height = dimensions['height']
            
            # 最小/最大値の設定
            scroll_width = max(scroll_width, 100)
            scroll_height = max(scroll_height, 100)
            scroll_width = min(scroll_width, 2000)
            scroll_height = min(scroll_height, 4000)
            
            # 高さに余白を追加
            adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
            adjusted_height = max(adjusted_height, scroll_height, 100)
            
            # ウィンドウサイズを調整
            driver.set_window_size(scroll_width, adjusted_height)
            time.sleep(0.2)
            
            # スクリーンショット取得
            logger.info("スクリーンショットを撮影...")
            png = driver.get_screenshot_as_png()
            
            # PIL画像に変換
            img = Image.open(BytesIO(png))
            logger.info(f"スクリーンショットサイズ: {img.width}x{img.height}")
            
            # 余白トリミング
            if trim_whitespace:
                img = trim_image_whitespace(img, threshold=248, padding=20)
                logger.info(f"トリミング後のサイズ: {img.width}x{img.height}")
            
            # 画像をHuggingFaceにアップロード
            prefix = f"infographic_{style}"
            image_url = hf_uploader.upload_image(img, prefix=prefix)
            
            elapsed = time.time() - start_time
            logger.info(f"並列処理による生成完了。所要時間: {elapsed:.2f}秒、URL: {image_url}")
            return img, image_url
            
        except Exception as e:
            logger.error(f"スクリーンショット生成中にエラー: {e}", exc_info=True)
            return Image.new('RGB', (1, 1), color=(0, 0, 0)), None
        finally:
            # WebDriverプールに戻す
            if driver_from_pool:
                driver_pool.release_driver(driver)
            # 一時ファイル削除
            if tmp_path and os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except Exception as e:
                    logger.error(f"一時ファイル削除エラー: {e}")
    
    except Exception as e:
        logger.error(f"並列処理中のエラー: {e}", exc_info=True)
        return Image.new('RGB', (1, 1), color=(0, 0, 0)), None  # エラー時は黒画像とNone URL

# 従来の非並列版も残す（互換性のため）
def text_to_screenshot(text: str, extension_percentage: float, temperature: float = 1.0,
                    trim_whitespace: bool = True, style: str = "standard") -> tuple:
    """テキストをGemini APIでHTMLに変換し、スクリーンショットを生成する統合関数（レガシー版）"""
    # 並列処理版を呼び出す
    return text_to_screenshot_parallel(text, extension_percentage, temperature, trim_whitespace, style)

# 新しい関数: HTMLからスクリーンショットを生成し、HuggingFaceにアップロード
def render_and_upload_screenshot(html_code: str, extension_percentage: float = 10.0, 
                                trim_whitespace: bool = True, prefix: str = "screenshot") -> tuple:
    """
    HTMLコードからスクリーンショットを生成し、HuggingFaceにアップロードする
    
    Returns:
        tuple - (PIL.Image, URL) - 生成された画像とHuggingFaceのURL
    """
    try:
        # スクリーンショット生成
        img = render_fullpage_screenshot(html_code, extension_percentage, trim_whitespace)
        
        # 画像をHuggingFaceにアップロード
        image_url = hf_uploader.upload_image(img, prefix=prefix)
        
        return img, image_url
    except Exception as e:
        logger.error(f"スクリーンショット生成とアップロード中にエラー: {e}", exc_info=True)
        return Image.new('RGB', (1, 1), color=(0, 0, 0)), None

# --- FastAPI Setup ---
app = FastAPI()

# CORS設定を追加
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静的ファイルのサービング設定
# Gradioのディレクトリを探索してアセットを見つける
gradio_dir = os.path.dirname(gr.__file__)
logger.info(f"Gradio version: {gr.__version__}")
logger.info(f"Gradio directory: {gradio_dir}")

# 基本的な静的ファイルディレクトリをマウント
static_dir = os.path.join(gradio_dir, "templates", "frontend", "static")
if os.path.exists(static_dir):
    logger.info(f"Mounting static directory: {static_dir}")
    app.mount("/static", StaticFiles(directory=static_dir), name="static")

# _appディレクトリを探す（新しいSvelteKitベースのフロントエンド用）
app_dir = os.path.join(gradio_dir, "templates", "frontend", "_app")
if os.path.exists(app_dir):
    logger.info(f"Mounting _app directory: {app_dir}")
    app.mount("/_app", StaticFiles(directory=app_dir), name="_app")

# assetsディレクトリを探す
assets_dir = os.path.join(gradio_dir, "templates", "frontend", "assets")
if os.path.exists(assets_dir):
    logger.info(f"Mounting assets directory: {assets_dir}")
    app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")

# cdnディレクトリがあれば追加
cdn_dir = os.path.join(gradio_dir, "templates", "cdn")
if os.path.exists(cdn_dir):
    logger.info(f"Mounting cdn directory: {cdn_dir}")
    app.mount("/cdn", StaticFiles(directory=cdn_dir), name="cdn")


# API Endpoint for screenshot generation - 更新版（URLを返すように変更）
@app.post("/api/screenshot",
          response_model=ImageUrlResponse,
          tags=["Screenshot"],
          summary="Render HTML to Full Page Screenshot and Upload to HuggingFace",
          description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, uploads to HuggingFace, and returns the URL.")
async def api_render_screenshot(request: ScreenshotRequest):
    """
    API endpoint to render HTML, upload to HuggingFace, and return the URL.
    """
    try:
        logger.info(f"API request received. Extension: {request.extension_percentage}%")
        
        # スクリーンショット生成とアップロード
        pil_image, image_url = render_and_upload_screenshot(
            request.html_code,
            request.extension_percentage,
            request.trim_whitespace,
            prefix="screenshot"
        )

        if pil_image.size == (1, 1) or not image_url:
             logger.error("Screenshot generation failed, or upload failed.")
             raise HTTPException(status_code=500, detail="Failed to generate or upload screenshot")

        # URLを返す
        logger.info(f"返却URL: {image_url}")
        return ImageUrlResponse(url=image_url)

    except Exception as e:
        logger.error(f"API Error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")

# --- 新しいGemini API連携エンドポイント（並列処理版）- URLを返すよう更新 ---
@app.post("/api/text-to-screenshot",
          response_model=ImageUrlResponse,
          tags=["Screenshot", "Gemini"],
          summary="テキストからインフォグラフィックを生成しHuggingFaceにアップロード",
          description="テキストをGemini APIを使ってHTMLインフォグラフィックに変換し、HuggingFaceにアップロードしたURLを返します。")
async def api_text_to_screenshot(request: GeminiRequest):
    """
    テキストからHTMLインフォグラフィックを生成してアップロードし、URLを返すAPIエンドポイント
    """
    try:
        logger.info(f"テキスト→スクリーンショットAPIリクエスト受信。テキスト長さ: {len(request.text)}, "
                   f"拡張率: {request.extension_percentage}%, 温度: {request.temperature}, "
                   f"スタイル: {request.style}")

        # 並列処理版を使用 - 画像とURLを取得
        pil_image, image_url = text_to_screenshot_parallel(
            request.text,
            request.extension_percentage,
            request.temperature,
            request.trim_whitespace,
            request.style
        )

        if pil_image.size == (1, 1) or not image_url:
            logger.error("スクリーンショット生成に失敗したか、アップロードに失敗しました。")
            raise HTTPException(status_code=500, detail="Failed to generate or upload screenshot")

        # URLを返す
        logger.info(f"返却URL: {image_url}")
        return ImageUrlResponse(url=image_url)

    except Exception as e:
        logger.error(f"API Error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")

# --- Gradio Interface Definition ---
# 入力モードの選択用Radioコンポーネント
def process_input(input_mode, input_text, extension_percentage, temperature, trim_whitespace, style):
    """入力モードに応じて適切な処理を行う"""
    if input_mode == "HTML入力":
        # HTMLモードの場合はレンダリングとアップロード
        img, url = render_and_upload_screenshot(
            input_text, 
            extension_percentage, 
            trim_whitespace, 
            prefix="html_screenshot"
        )
        return img, url if url else "アップロード失敗またはURL取得できませんでした"
    else:
        # テキスト入力モードの場合はGemini API使用（並列処理版）
        img, url = text_to_screenshot_parallel(
            input_text, 
            extension_percentage, 
            temperature, 
            trim_whitespace, 
            style
        )
        return img, url if url else "アップロード失敗またはURL取得できませんでした"

# Gradio UIの定義
with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Base()) as iface:
    gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック変換")
    gr.Markdown("HTMLコードをレンダリングするか、テキストをGemini APIでインフォグラフィックに変換して画像として取得します。")
    gr.Markdown("**パフォーマンス向上版**: 並列処理と最適化により処理時間を短縮しています")

    with gr.Row():
        input_mode = gr.Radio(
            ["HTML入力", "テキスト入力"],
            label="入力モード",
            value="HTML入力"
        )

    # 共用のテキストボックス
    input_text = gr.Textbox(
        lines=15,
        label="入力",
        placeholder="HTMLコードまたはテキストを入力してください。入力モードに応じて処理されます。"
    )

    with gr.Row():
        with gr.Column(scale=1):
            # スタイル選択ドロップダウン
            style_dropdown = gr.Dropdown(
                choices=["standard", "cute", "resort", "cool", "dental", "school","KOKUGO"],
                value="standard",
                label="デザインスタイル",
                info="テキスト→HTML変換時のデザインテーマを選択します",
                visible=False  # テキスト入力モードの時だけ表示
            )

        with gr.Column(scale=2):
            extension_percentage = gr.Slider(
                minimum=0,
                maximum=30,
                step=1.0,
                value=15,    # デフォルト値15%
                label="上下高さ拡張率（%）"
            )

            # 温度調整スライダー（テキストモード時のみ表示）
            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=1.0,    # デフォルト値を1.0に設定
                label="生成時の温度（低い=一貫性高、高い=創造性高）",
                visible=False  # 最初は非表示
            )

    # 余白トリミングオプション
    trim_whitespace = gr.Checkbox(
        label="余白を自動トリミング",
        value=True,
        info="生成される画像から余分な空白領域を自動的に削除します"
    )

    submit_btn = gr.Button("生成")
    
    # 出力部分をRowで分ける
    with gr.Row():
        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="ページ全体のスクリーンショット")
        with gr.Column(scale=1):
            output_url = gr.Textbox(label="画像URL（HuggingFace）", info="生成された画像のURLです。このURLを使用して画像にアクセスできます。")

    # 入力モード変更時のイベント処理（テキストモード時のみ温度スライダーとスタイルドロップダウンを表示）
    def update_controls_visibility(mode):
        # Gradio 4.x用のアップデート方法
        is_text_mode = mode == "テキスト入力"
        return [
            gr.update(visible=is_text_mode),  # temperature
            gr.update(visible=is_text_mode),  # style_dropdown
        ]

    input_mode.change(
        fn=update_controls_visibility,
        inputs=input_mode,
        outputs=[temperature, style_dropdown]
    )

    # 生成ボタンクリック時のイベント処理
    submit_btn.click(
        fn=process_input,
        inputs=[input_mode, input_text, extension_percentage, temperature, trim_whitespace, style_dropdown],
        outputs=[output_image, output_url]
    )

    # 環境変数情報を表示
    gemini_model = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
    hf_repo = os.environ.get("HF_REPO_ID", "leave-everything/ChoTensaiJinGrareko")
    gr.Markdown(f"""
    ## APIエンドポイント
    - `/api/screenshot` - HTMLコードからスクリーンショットを生成し、URLを返します
    - `/api/text-to-screenshot` - テキストからインフォグラフィックスクリーンショットを生成し、URLを返します

    ## 設定情報
    - 使用モデル: {gemini_model} (環境変数 GEMINI_MODEL で変更可能)
    - HuggingFaceリポジトリ: {hf_repo} (環境変数 HF_REPO_ID で変更可能)
    - WebDriverプール最大数: {driver_pool.max_drivers} (環境変数 MAX_WEBDRIVERS で変更可能)
    - 対応スタイル: standard, cute, resort, cool, dental, KOKUGO
    """)

# --- Mount Gradio App onto FastAPI ---
app = gr.mount_gradio_app(app, iface, path="/")

# --- Run with Uvicorn (for local testing) ---
if __name__ == "__main__":
    import uvicorn
    logger.info("Starting Uvicorn server for local development...")
    uvicorn.run(app, host="0.0.0.0", port=7860)

# アプリケーション終了時にWebDriverプールをクリーンアップ
import atexit
atexit.register(driver_pool.close_all)