app.py

import os
import tempfile
import gradio as gr
import pandas as pd
from urllib.parse import urlparse, parse_qs
from PIL import Image
from io import BytesIO
import requests
from google import genai

# 請確保你的 requirements.txt 包含 pandas
# pip install pandas

# --- Google Sheets 相關函式 (已更新) ---
def read_google_sheet(sheet_url: str):
    """
    從 Google Sheet 的 URL 讀取資料。
    """
    if not sheet_url:
        raise gr.Error("請提供 Google Sheet URL。")
    try:
        def build_csv_url(url: str) -> str:
            parsed = urlparse(url)
            doc_id = parsed.path.strip("/").split("/")[2] if len(parsed.path.strip("/").split("/")) >= 3 and parsed.path.strip("/").split("/")[1] == "d" else None
            gid = parse_qs(parsed.query).get("gid", [None])[0] or parse_qs(parsed.fragment).get("gid", [None])[0] or "0"
            if doc_id:
                return f"https://docs.google.com/spreadsheets/d/{doc_id}/export?format=csv&gid={gid}"
            if "/export" in parsed.path and "format=csv" in parsed.query:
                return url
            return url.replace("/edit#gid=0", "/export?format=csv&gid=0")

        csv_url = build_csv_url(sheet_url)
        df = pd.read_csv(csv_url, engine='python', on_bad_lines='warn', encoding='utf-8')
        return df
    except Exception as e:
        raise gr.Error(f"讀取 Google Sheet 時發生錯誤: {e}")

def process_sheet_data(sheet_url):
    """處理試算表資料，為 Gradio DataFrame 準備。"""
    try:
        df = read_google_sheet(sheet_url)
        if df.shape[1] < 4:  # 檢查是否至少有 4 列 (索引, 白背圖URL, 參考圖URL, 提示詞)
            error_msg = f"錯誤：Google Sheet 至少需要 4 列 (索引, 白背圖URL, 參考圖URL, 提示詞)。目前只有 {df.shape[1]} 列。"
            raise gr.Error(error_msg)
        
        # 假設列順序為: White_Back_Image_URL, Ref_Image_URL, Prompt
        headers = ["Index", "白背圖URL", "參考圖URL", "提示詞"]
        data_list = []
        for i, row in df.iterrows():
            if pd.notna(row.iloc[0]):  # 檢查第一列 (白背圖URL) 是否有資料
                data_list.append([i + 2, row.iloc[0], row.iloc[1], row.iloc[2]])
        
        log_message = f"成功讀取 {len(data_list)} 筆數據。"
        return data_list, log_message
    except Exception as e:
        raise gr.Error(f"處理試算表資料時發生錯誤: {e}")

def get_row_data(sheet_url, row_number):
    """從 Google Sheet 讀取指定列的資料。"""
    try:
        df = read_google_sheet(sheet_url)
        row_index = int(row_number) - 2 # 由於DataFrame索引從0開始，而Gradio顯示的行數從2開始
        if row_index < 0 or row_index >= df.shape[0]:
            raise gr.Error(f"指定的行數 {row_number} 不存在。")
        
        row_data = df.iloc[row_index]
        white_back_image_url = row_data.iloc[0] if pd.notna(row_data.iloc[0]) else ""
        ref_image_url = row_data.iloc[1] if pd.notna(row_data.iloc[1]) else ""
        prompt_text = row_data.iloc[2] if pd.notna(row_data.iloc[2]) else ""
        
        return white_back_image_url, ref_image_url, prompt_text
    except Exception as e:
        raise gr.Error(f"讀取指定行時發生錯誤: {e}")

# --- 下載圖片函式 ---
def load_image_from_url(url: str):
    """從 URL 下載圖片並以 PIL Image 格式回傳。"""
    if not url or not isinstance(url, str):
        return None
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        }
        response = requests.get(url, timeout=20, headers=headers)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content)).convert("RGB")
        return image
    except requests.exceptions.HTTPError as e:
        gr.Warning(f"下載圖片失敗：HTTP 錯誤 {e.response.status_code}")
        return None
    except Exception as e:
        gr.Warning(f"下載圖片時發生意外錯誤：{e}")
        return None

# --- Gemini 核心函式 ---
def generate_image(text, images, api_key, model="gemini-2.5-flash-image-preview"):
    """使用 Gemini 模型生成圖片。"""
    if not api_key or api_key.strip() == "":
        raise gr.Error("請輸入有效的 Gemini API 金鑰。", duration=10)
    try:
        client = genai.Client(api_key=api_key.strip())
        contents = images + [text]
        response = client.models.generate_content(model=model, contents=contents)
        text_response = ""
        image_path = None
        for part in response.candidates[0].content.parts:
            if part.text is not None:
                text_response += part.text + "\n"
            elif part.inline_data is not None:
                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                    temp_path = tmp.name
                    generated_image = Image.open(BytesIO(part.inline_data.data))
                    generated_image.save(temp_path)
                    image_path = temp_path
        return image_path, text_response
    except Exception as e:
        raise gr.Error(f"Gemini API 呼叫失敗: {e}", duration=10)

# --- Gradio 互動函式 ---
def generate_image_from_row(sheet_url, row_number, gemini_api_key):
    """根據指定的行數，生成圖片。"""
    if not sheet_url:
        raise gr.Error("請先輸入 Google Sheet URL。", duration=5)
    if not row_number or row_number <= 1:
        raise gr.Error("請輸入有效的行數 (大於 1)。", duration=5)
        
    try:
        white_back_url, ref_image_url, prompt_text = get_row_data(sheet_url, row_number)
        
        log_message = f"開始處理第 {row_number} 行...\n"
        log_message += f"白背圖URL: {white_back_url}\n"
        log_message += f"參考圖URL: {ref_image_url}\n"
        log_message += f"提示詞: {prompt_text}\n"

        images_for_gemini = []
        if white_back_url:
            wb_img = load_image_from_url(white_back_url)
            if wb_img:
                images_for_gemini.append(wb_img)
        if ref_image_url:
            ref_img = load_image_from_url(ref_image_url)
            if ref_img:
                images_for_gemini.append(ref_img)

        if not images_for_gemini:
            return None, "警告：無效的圖片 URL，無法生成圖片。"
        
        log_message += "圖片已下載，開始呼叫 Gemini 模型...\n"
        
        image_path, text_response = generate_image(text=prompt_text, images=images_for_gemini, api_key=gemini_api_key)

        if image_path:
            log_message += "圖片生成成功！"
            return image_path, log_message
        else:
            log_message += "圖片生成失敗。\n"
            log_message += f"Gemini 文字回應: {text_response}"
            return None, log_message
    except ValueError:
        raise gr.Error("行數必須為數字。", duration=5)
    except Exception as e:
        return None, f"生成圖片時發生錯誤: {e}"

# --- Gradio 介面設定 ---
with gr.Blocks() as demo:
    gr.Markdown("## Google Sheets 圖片生成器")
    
    with gr.Row(elem_classes="main-content"):
        with gr.Column(elem_classes="input-column"):
            gemini_api_key = gr.Textbox(
                lines=1,
                placeholder="請在此輸入你的 Gemini API 金鑰",
                label="Gemini API 金鑰",
                elem_classes="api-key-input"
            )
            sheet_url_input = gr.Textbox(
                label="Google Sheet URL", 
                value="https://docs.google.com/spreadsheets/d/1G3olHxydDIbnyXdh5nnw5TG0akZFeMeYm-25JmCGDLg/edit?gid=0#gid=0"
            )
            process_button = gr.Button("處理試算表", elem_classes="generate-btn")
            
            with gr.Row():
                row_index_input = gr.Number(label="要生成的行數", precision=0, value=2)
                generate_selected_button = gr.Button("生成所選行的圖片", elem_classes="generate-btn")
        
        with gr.Column(elem_classes="output-column"):
            output_dataframe = gr.DataFrame(
                headers=["Index", "白背圖URL", "參考圖URL", "提示詞"],
                col_count=(4, "fixed"),
                interactive=False,
                label="已處理的試算表數據"
            )
            generated_image_output = gr.Image(label="生成的圖片", elem_classes="output-gallery")
            operation_log_output = gr.Textbox(
                label="操作日誌", 
                lines=10,
                placeholder="文字回應和日誌會顯示在這裡。"
            )

    # 按鈕的事件綁定
    process_button.click(
        fn=process_sheet_data,
        inputs=[sheet_url_input],
        outputs=[output_dataframe, operation_log_output]
    )

    generate_selected_button.click(
        fn=generate_image_from_row,
        inputs=[sheet_url_input, row_index_input, gemini_api_key],
        outputs=[generated_image_output, operation_log_output]
    )
    
    demo.queue().launch(mcp_server=True, share=True)