import gradio as gr
from google import genai
import pandas as pd

def fetch_gemini_models(api_key):
    """
    Fetches available Gemini models and returns them as a pandas DataFrame.
    """



    if not api_key:
        return None, "エラー: GEMINI_API_KEY が必要です。"

    try:
        client = genai.Client(api_key=api_key)
        models = list(client.models.list())
        models.sort(key=lambda x: x.name)

        # Mapping for Japanese translation
        method_translation = {
            "generateContent": "コンテンツ生成",
            "countTokens": "トークン計算",
            "createCachedContent": "キャッシュ作成",
            "embedText": "テキスト埋め込み",
            "batchGenerate": "バッチ生成",
            "embedContent": "コンテンツ埋め込み"
        }

        data = []
        for m in models:
            if hasattr(m, "supported_actions") and m.supported_actions and "generateContent" in m.supported_actions:
                translated_methods = [method_translation.get(method, method) for method in m.supported_actions]
                methods = ", ".join(translated_methods)
                
                # Checkbox column first
                data.append({
                    "選択": False,
                    "モデル名": m.name,
                    "表示名": m.display_name,
                    "対応メソッド": methods
                })
        
        if not data:
            return None, "'generateContent' をサポートするモデルが見つかりませんでした。"
        
        df = pd.DataFrame(data)
        return df, "成功: モデルの取得に成功しました。"
            

    except Exception as e:
        return None, f"エラーが発生しました: {str(e)}"

def parse_error_message(error_str):
    """
    Parses complex error messages and returns a concise Japanese explanation.
    """
    if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:
        return "リソース制限を超過しました (429 Resource Exhausted) - 無料枠の上限に達した可能性があります。"
    elif "403" in error_str or "PERMISSION_DENIED" in error_str:
        return "アクセス権限がありません (403 Permission Denied) - APIキーの設定やモデルの利用権限を確認してください。"
    elif "404" in error_str or "NOT_FOUND" in error_str:
        return "モデルが見つかりません (404 Not Found) - 指定されたモデルが存在しないか、利用できません。"
    elif "400" in error_str or "INVALID_ARGUMENT" in error_str:
        return "無効なリクエストです (400 Invalid Argument) - パラメータ設定を確認してください。"
    else:
        # Return a shortened version of the original error if it's too long
        return f"エラー: {error_str[:100]}..." if len(error_str) > 100 else f"エラー: {error_str}"

def check_selected_models(api_key, df):
    """
    Checks access for all selected models in the dataframe.
    """
    if not api_key:
        return "エラー: GEMINI_API_KEY が必要です。"
    if df is None or df.empty:
        return "エラー: モデル一覧が空です。"

    # Filter selected rows
    # Gradio dataframe usually sends back a pandas dataframe or list of lists
    if isinstance(df, dict): # Handle some Gradio versions wrapping it
        df = pd.DataFrame(df['data'], columns=df['headers'])
    
    # Check if "選択" column exists
    if "選択" not in df.columns:
        return "エラー: モデル一覧の形式が不正です。"

    selected_rows = df[df["選択"] == True]
    
    if selected_rows.empty:
        return "警告: チェックするモデルが選択されていません。"

    results = []
    client = genai.Client(api_key=api_key)

    for index, row in selected_rows.iterrows():
        model_name = row["モデル名"]
        try:
            # Try a minimal generation request
            response = client.models.generate_content(
                model=model_name,
                contents="test",
            )
            results.append(f"✅ {model_name}: アクセス確認成功")
        except Exception as e:
            error_msg = parse_error_message(str(e))
            results.append(f"❌ {model_name}: 失敗 - {error_msg}")
    
    return "\n".join(results)

# Define Gradio Theme (Rich Aesthetics)
theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
).set(
    button_primary_background_fill='*primary_600',
    button_primary_background_fill_hover='*primary_700',
)

with gr.Blocks(theme=theme, title="Gemini モデルチェッカー") as demo:
    gr.Markdown(
        """
        # 💎 Gemini モデルチェッカー
        Hugging Face Spaces上のGradioで、利用可能なGeminiモデルを簡単に確認できます。
        """
    )
    
    with gr.Row():
        api_key_input = gr.Textbox(
            label="Gemini APIキー",
            placeholder="AIzaSy...",
            type="password",
            interactive=True,
            scale=4
        )
        fetch_btn = gr.Button("モデル一覧取得", variant="primary", scale=1)
    
    output_status = gr.Textbox(label="ステータス", interactive=False)
    # Make dataframe interactive to allow checkbox editing
    output_table = gr.DataFrame(
        label="利用可能なモデル (左端のチェックボックスで選択)",
        interactive=True,
        type="pandas",
        col_count=(4, "fixed"), # 選択, モデル名, 表示名, 対応メソッド
        headers=["選択", "モデル名", "表示名", "対応メソッド"],
        datatype=["bool", "str", "str", "str"]
    )

    fetch_btn.click(
        fn=fetch_gemini_models,
        inputs=api_key_input,
        outputs=[output_table, output_status]
    )

    gr.Markdown("## モデルアクセス確認")
    gr.Markdown("上の表でチェックを入れたモデルに対して、一括でアクセス確認を行います。")
    
    check_btn = gr.Button("選択したモデルのアクセス確認", variant="secondary")
    access_status = gr.Textbox(label="アクセス確認結果", interactive=False, lines=10)

    check_btn.click(
        fn=check_selected_models,
        inputs=[api_key_input, output_table],
        outputs=access_status
    )

    gr.Markdown(
        """
        ---
        ### 使い方
        1. Google AI Studioから取得した **APIキー** を入力します。
        2. **モデル一覧取得** ボタンをクリックします。利用可能なモデルの一覧が表示されます。
        3. 特定のモデルへのアクセスを確認するには、**モデルアクセス確認** セクションで **モデル名** を選択（または入力）し、**アクセス確認** ボタンをクリックします。
        """
    )

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