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

# --- 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)
        # 增加 read_csv 的 robustness，防止某些 Google Sheet 導致的 ParserError
        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)
        # 檢查是否至少有 3 列 (白背圖URL, 參考圖URL, 提示詞)，因為索引會在程式中生成
        if df.shape[1] < 3:
            error_msg = f"錯誤：Google Sheet 至少需要 3 列 (白背圖URL, 參考圖URL, 提示詞)。目前只有 {df.shape[1]} 列。"
            raise gr.Error(error_msg)
        
        data_list = []
        for i, row in df.iterrows():
            # 確保有足夠的列且至少第一列(白背圖URL)不為空才加入
            if len(row) >= 3 and pd.notna(row.iloc[0]):
                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)
        # Gradio顯示的行數從2開始，DataFrame索引從0開始，所以需要減2
        row_index = int(row_number) - 2 
        if row_index < 0 or row_index >= df.shape[0]:
            raise gr.Error(f"指定的行數 {row_number} 不存在。")
        
        row_data = df.iloc[row_index]
        # 確保取值時索引不越界，並處理 NaN
        white_back_image_url = row_data.iloc[0] if len(row_data) > 0 and pd.notna(row_data.iloc[0]) else ""
        ref_image_url = row_data.iloc[1] if len(row_data) > 1 and pd.notna(row_data.iloc[1]) else ""
        prompt_text = row_data.iloc[2] if len(row_data) > 2 and 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) or not url.strip(): # 增加對空字串的檢查
        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 if white_back_url else '無'}\n"
        log_message += f"參考圖URL: {ref_image_url if ref_image_url else '無'}\n"
        log_message += f"提示詞: {prompt_text if prompt_text else '無'}\n"

        images_for_gemini = []
        
        # 下載圖片
        wb_img = None
        if white_back_url:
            wb_img = load_image_from_url(white_back_url)
            if wb_img:
                images_for_gemini.append(wb_img)
            else:
                log_message += f"警告：無法下載白背圖 '{white_back_url}'。\n"

        ref_img = None
        if ref_image_url:
            ref_img = load_image_from_url(ref_image_url)
            if ref_img:
                images_for_gemini.append(ref_img)
            else:
                log_message += f"警告：無法下載參考圖 '{ref_image_url}'。\n"

        if not images_for_gemini:
            # 如果沒有任何圖片，但有提示詞，嘗試純文字生成
            if prompt_text:
                log_message += "沒有可用的輸入圖片，嘗試僅使用提示詞生成圖片。\n"
            else:
                return None, None, None, "警告：無效的圖片 URL 且無提示詞，無法生成圖片。", log_message

        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 wb_img, ref_img, image_path, log_message
        else:
            log_message += "圖片生成失敗。\n"
            log_message += f"Gemini 文字回應: {text_response}"
            return wb_img, ref_img, None, log_message
            
    except ValueError:
        return None, None, None, "行數必須為數字。", "生成圖片時發生錯誤：行數必須為數字。"
    except Exception as e:
        return None, None, None, f"生成圖片時發生錯誤: {e}", 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",
                # 從環境變數讀取預設值 (如果存在)
                value=os.environ.get("GEMINI_API_KEY", "") 
            )
            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="已處理的試算表數據"
            )
            
            # 新增顯示白背圖和參考圖的欄位
            with gr.Row():
                original_white_back_output = gr.Image(label="原始白背圖", elem_classes="output-gallery", interactive=False)
                original_ref_image_output = gr.Image(label="原始參考圖", elem_classes="output-gallery", interactive=False)
            
            generated_image_output = gr.Image(label="生成的圖片", elem_classes="output-gallery", interactive=False)
            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=[original_white_back_output, original_ref_image_output, generated_image_output, operation_log_output]
    )
    
    demo.queue().launch(mcp_server=True, share=True)