import gradio as gr
import numpy as np
import random
from PIL import Image
from io import BytesIO
import os

# Google Generative AIライブラリをインポート
from google import genai
from google.genai import types


def infer(
    google_api_key,
    prompt,
    progress=gr.Progress(track_tqdm=True),
):
    # Google APIキーが入力されているかチェック
    if not google_api_key:
        raise gr.Error("Google API Key is required. Please enter your key.")

    try:
        # Clientの初期化時にAPIキーを渡す
        client = genai.Client(api_key=google_api_key)
        
        # プロンプトにシード情報を含めることで、再現性のヒントとする
        full_prompt = f"{prompt}"

        response = client.models.generate_content(
            model="gemini-2.5-flash-image-preview",
            contents=[full_prompt],
        )

        # レスポンスから画像データを抽出
        image_part = None
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image_part = part
                break

        if image_part:
            image = Image.open(BytesIO(image_part.inline_data.data))
            return image
        else:
            # テキスト部分があればエラーとして表示
            text_response = ""
            for part in response.candidates[0].content.parts:
                if part.text is not None:
                    text_response += part.text
            raise gr.Error(f"Image could not be generated. API Response: {text_response or 'No content found'}")

    except Exception as e:
        raise gr.Error(f"An error occurred during image generation: {e}")

examples = [
    "Create a picture of a nano banana dish in a fancy restaurant with a Gemini theme",
    "An astronaut riding a green horse in a photorealistic style",
    "A delicious ceviche cheesecake slice, studio lighting",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Image with Google Gemini")

        google_api_key = gr.Textbox(
            label="Google API Key",
            placeholder="Enter your Google API Key here...",
            type="password",
            show_label=True,
        )

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0, variant="primary")

        result = gr.Image(label="Result", show_label=False)

        # --- Advanced Settingsを削除 ---

        gr.Examples(examples=examples, inputs=[prompt])

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            google_api_key,
            prompt,
        ],
        outputs=[result],
    )

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