# app.py

import gradio as gr
#from together import Together
import together

import os
import base64
from PIL import Image
from io import BytesIO

# --- 1. 設定: APIキーの読み込み ---
api_key = os.environ.get("TOGETHER_API_KEY")
if api_key:
    together_client = together.Together(api_key=api_key)
else:
    together_client = None

# 使用する画像生成モデル
IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell-Free"

# --- 2. 中核機能: 画像生成関数 ---
def generate_image_from_prompt(prompt: str, history: list):
    """
    TogetherAI APIを使用して画像を生成し、履歴リストの先頭に追加します。
    """
    if not together_client:
        raise gr.Error("TogetherAIのAPIクライアントが設定されていません。環境変数 TOGETHER_API_KEY を確認してください。")
    if not prompt or not prompt.strip():
        raise gr.Error("画像を生成するには、プロンプトを入力してください。")

    print(f"プロンプトを元に画像を生成中: '{prompt}'")
    try:
        # TogetherAI APIを呼び出し
        response = together_client.images.generate(
            prompt=prompt,
            model=IMAGE_MODEL,
            steps=4,
            n=4,
            height=1024,
            width=1024
        )

         # ▼▼▼ デバッグのためにこの行を追加 ▼▼▼
        print("---------- API Response ----------")
        print(response)
        print("--------------------------------")
        
        generated_images = []
        # response.data にはリクエストした全ての画像データが含まれる
        for image_data_item in response.data:
            image_b64 = image_data_item.b64_json
            
            # ▼▼▼ 修正箇所: 安全チェックを追加 ▼▼▼
            # image_b64がNoneでない（正常にデータが返された）場合のみ処理する
            if image_b64:
                image_data = base64.b64decode(image_b64)
                image = Image.open(BytesIO(image_data))
                # Gallery用に (image, caption) のタプルを作成
                generated_images.append((image, prompt))

        # 正常に生成された画像が1枚もなかった場合
        if not generated_images:
            raise gr.Error("画像の生成に失敗しました。プロンプトがコンテンツフィルターに抵触した可能性があります。別のプロンプトで試してください。")

        # 新しい画像のリストを、既存の履歴の先頭に追加
        new_history = generated_images + history
        
        print(f"{len(generated_images)}枚の画像の生成が完了しました。")
        return new_history, new_history

    except Exception as e:
        # 既にgr.Errorの場合はそのまま表示
        if isinstance(e, gr.Error):
            raise e
        print(f"予期せぬエラーが発生しました: {e}")
        raise gr.Error(f"画像の生成に失敗しました。エラー詳細: {e}")

# --- 3. Gradio UIの構築 ---
with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 800px; margin: auto;}") as app:
    gr.Markdown(
        """
        # 🖼️ TogetherAI 画像生成アプリ (Gradio版)
        プロンプトを入力して、TogetherAIのAPIで画像を生成します。
        最新の画像がギャラリーの先頭に表示されます。
        """
    )
    gr.Markdown("---")
    
    history_state = gr.State([])

    with gr.Column():
        prompt_input = gr.Textbox(
            label="画像生成プロンプト",
            placeholder="例: 月面にいる、写真のようにリアルな猫の宇宙飛行士",
            lines=4,
            interactive=True,
        )
        if not together_client:
             gr.Warning("環境変数 `TOGETHER_API_KEY` が見つかりません。画像生成機能を利用するにはAPIキーを設定してください。")
        generate_btn = gr.Button(
            "画像生成 ✨",
            variant="primary",
            interactive=(together_client is not None)
        )

    gr.Markdown("---")
    gr.Markdown("## 生成された画像")
    
    image_gallery = gr.Gallery(
        label="生成結果",
        show_label=False,
        columns=2,
        height="auto",
        object_fit="contain",
        value=None,
        preview=True
    )
    
    generate_btn.click(
        fn=generate_image_from_prompt,
        inputs=[prompt_input, history_state],
        outputs=[image_gallery, history_state],
        show_progress="full"
    )
    
    gr.Examples(
        examples=[
            "A high-resolution photo of a futuristic city with flying cars",
            "An oil painting of a serene landscape with a river and mountains",
            "A cute, fluffy alien creature exploring a vibrant, magical forest",
            "Logo for a coffee shop named 'The Daily Grind', minimalist, modern"
        ],
        inputs=prompt_input
    )

if __name__ == "__main__":
    app.launch(mcp_server=True)