# app.py

import gradio as gr
import together
import os
import base64
from PIL import Image
from io import BytesIO
import requests  # ▼▼▼ 修正箇所: requestsライブラリをインポート ▼▼▼

# --- 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):
    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:
        response = together_client.images.generate(
            prompt=prompt,
            model=IMAGE_MODEL,
            steps=4,
            n=4,
            height=1024,
            width=1024
        )
        
        generated_images = []
        
        # ▼▼▼ 修正箇所: URLから画像をダウンロードするロジックに変更 ▼▼▼
        for image_data_item in response.data:
            # APIの応答に 'url' が含まれているかチェック
            if hasattr(image_data_item, 'url') and image_data_item.url:
                try:
                    image_url = image_data_item.url
                    print(f"URLから画像をダウンロード中: {image_url}")
                    
                    # URLにアクセスして画像データを取得
                    image_response = requests.get(image_url)
                    image_response.raise_for_status()  # エラーがあれば例外を発生させる
                    
                    # ダウンロードしたバイトデータからPILイメージを作成
                    image = Image.open(BytesIO(image_response.content))
                    generated_images.append((image, prompt))
                    
                except requests.exceptions.RequestException as e:
                    print(f"画像のダウンロードに失敗しました: {e}")
                    continue # この画像の処理をスキップして次に進む
            
            # (参考) b64_jsonが返ってくるモデルのためのフォールバック処理
            elif image_data_item.b64_json:
                image_b64 = image_data_item.b64_json
                image_data = base64.b64decode(image_b64)
                image = Image.open(BytesIO(image_data))
                generated_images.append((image, prompt))

        if not generated_images:
            raise gr.Error("画像の生成に失敗しました。APIから有効な画像データやURLが返されませんでした。")

        new_history = generated_images + history
        print(f"{len(generated_images)}枚の画像の生成が完了しました。")
        return new_history, new_history

    except Exception as e:
        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("画像生成 ✨ (4枚)", 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)