main.py

from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from linebot import LineBotApi, WebhookHandler
from linebot.exceptions import InvalidSignatureError
from linebot.models import MessageEvent, TextMessage, TextSendMessage, ImageMessage, ImageSendMessage
import json
import os
import requests
import base64
from collections import defaultdict
import uvicorn
import logging
from PIL import Image
import io

# 設置日誌記錄
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 檢查環境變數
if not all(k in os.environ for k in ["CHANNEL_ACCESS_TOKEN", "CHANNEL_SECRET", "GOOGLE_API_KEY", "IMGBB_API_KEY"]):
    logging.error("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and IMGBB_API_KEY.")
    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and IMGBB_API_KEY.")

# Line Bot API 設定
line_bot_api = LineBotApi(os.environ["CHANNEL_ACCESS_TOKEN"])
line_handler = WebhookHandler(os.environ["CHANNEL_SECRET"])

# 使用者狀態追蹤
user_states = defaultdict(lambda: {
    "upper_body_images": [],
    "lower_body_images": [],
    "current_mode": None, # "upper" or "lower"
    "is_ready_for_outfit": False, # 標記衣物是否收集完畢
    "is_ready_for_photo": False, # 標記個人照片是否上傳
    "user_info": {}, # 儲存身高、三圍、場合
    "personal_photo": None, # 儲存個人照片
    "personal_photo_base64": None # 新增：儲存個人照片的 Base64 字串
})

MAX_IMAGES_PER_TYPE = 3
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-05-20:generateContent"
IMAGIN_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent"
GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
IMGBB_API_KEY = os.environ["IMGBB_API_KEY"]

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def root():
    return {"title": "Line Bot 穿搭建議"}

@app.post("/webhook")
async def webhook(
    request: Request,
    background_tasks: BackgroundTasks,
    x_line_signature=Header(None),
):
    body = await request.body()
    try:
        background_tasks.add_task(line_handler.handle, body.decode("utf-8"), x_line_signature)
    except InvalidSignatureError:
        raise HTTPException(status_code=400, detail="Invalid signature")
    return "ok"

def get_base64_image(message_id: str):
    logging.info(f"Fetching image content for message ID: {message_id}")
    message_content = line_bot_api.get_message_content(message_id)
    image_bytes = message_content.content
    return base64.b64encode(image_bytes).decode('utf-8')

def upload_to_imgbb(image_base64: str):
    """
    將 Base64 編碼的圖片上傳到 Imgbb。
    """
    logging.info("Attempting to upload image to Imgbb.")
    url = "https://api.imgbb.com/1/upload"
    payload = {
        "key": IMGBB_API_KEY,
        "image": image_base64
    }
    try:
        response = requests.post(url, data=payload)
        response.raise_for_status()
        imgbb_link = response.json()["data"]["url"]
        logging.info(f"Image successfully uploaded to Imgbb: {imgbb_link}")
        return imgbb_link
    except requests.exceptions.RequestException as e:
        logging.error(f"Error uploading to Imgbb: {e}")
        if 'response' in locals() and response is not None:
            logging.error(f"Imgbb API Response: {response.status_code} - {response.text}")
        return None

def get_base64_from_url(image_url: str):
    """
    從 URL 下載圖片並轉換為 Base64 編碼。
    """
    try:
        response = requests.get(image_url)
        response.raise_for_status()
        return base64.b64encode(response.content).decode('utf-8')
    except requests.exceptions.RequestException as e:
        logging.error(f"Error fetching image from URL: {image_url}, error: {e}")
        return None

def get_gemini_response(prompt: str, images: list, api_url: str):
    logging.info("Sending request to Gemini API.")
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ] + [
                    {"inlineData": {"mimeType": "image/jpeg", "data": img}} for img in images
                ]
            }
        ]
    }
    
    try:
        response = requests.post(f"{api_url}?key={GOOGLE_API_KEY}", json=payload, timeout=120) # Increased timeout for image generation
        response.raise_for_status()
        response_data = response.json()
        if 'candidates' in response_data and response_data['candidates']:
            gemini_text = response_data['candidates'][0]['content']['parts'][0]['text']
            logging.info("Successfully received response from Gemini.")
            return gemini_text
        else:
            logging.error(f"Gemini API response has no candidates: {response.text}")
            return "Gemini API 暫時無法提供服務，請稍後再試。"
    except requests.exceptions.RequestException as e:
        logging.error(f"Gemini API request failed: {e}")
        return f"Gemini API 請求失敗：{e}"


def get_virtual_tryon_image(user_photo_base64: str, upper_body_image_base64: str, lower_body_image_base64: str):
    """
    使用 gemini-2.5-flash-image-preview (nanobanana) 模擬虛擬試穿。
    """
    logging.info("Attempting virtual try-on with gemini-2.5-flash-image-preview.")
    
    prompt = "請將提供的上衣和褲子，虛擬試穿到第一張人像照片上。結果必須看起來真實且自然，衣服的紋理和細節必須保留。"
    
    all_images_base64 = [user_photo_base64, upper_body_image_base64, lower_body_image_base64]
    
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ] + [
                    {"inlineData": {"mimeType": "image/jpeg", "data": img_base64}} for img_base64 in all_images_base64
                ]
            }
        ],
        "generationConfig": {
            "responseModalities": ['IMAGE']
        }
    }
    
    try:
        response = requests.post(f"{IMAGIN_API_URL}?key={GOOGLE_API_KEY}", json=payload, timeout=120)
        response.raise_for_status()
        response_data = response.json()
        if 'candidates' in response_data and response_data['candidates']:
            # 從回應中提取 Base64 編碼的圖片
            image_part = response_data['candidates'][0]['content']['parts'][0]
            if image_part and 'inlineData' in image_part:
                generated_image_base64 = image_part['inlineData']['data']
                # 將生成的 Base64 圖片上傳到 Imgbb
                virtual_tryon_url = upload_to_imgbb(generated_image_base64)
                if virtual_tryon_url:
                    logging.info(f"Successfully generated and uploaded virtual try-on image: {virtual_tryon_url}")
                    return virtual_tryon_url
                else:
                    logging.error("Failed to upload generated image to Imgbb.")
                    return None
            else:
                logging.error("Gemini image response format is invalid.")
                return None
        else:
            logging.error(f"Gemini image generation response has no candidates: {response.text}")
            return None
    except requests.exceptions.RequestException as e:
        logging.error(f"Gemini image generation API request failed: {e}")
        return None

@line_handler.add(MessageEvent, message=TextMessage)
def handle_text_message(event):
    user_id = event.source.user_id
    text = event.message.text.lower()
    reply_token = event.reply_token

    logging.info(f"Received text message from user {user_id}: '{text}'")

    # 新增 重置 功能
    if text in ["重置", "重來", "重新開始", "再一次"]:
        user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None, "is_ready_for_outfit": False, "is_ready_for_photo": False, "user_info": {}, "personal_photo": None, "personal_photo_base64": None})[user_id]
        line_bot_api.reply_message(
            reply_token,
            TextSendMessage(text="狀態已重置。請重新輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
        )
        return
    
    # 處理個人資訊輸入
    if not user_states[user_id]["is_ready_for_photo"] and not user_states[user_id]["is_ready_for_outfit"]:
        try:
            parts = text.split(',')
            if len(parts) == 5:
                height = float(parts[0])
                bust = float(parts[1])
                waist = float(parts[2])
                hip = float(parts[3])
                occasion = parts[4].strip()
                
                user_states[user_id]["user_info"] = {
                    "height": height,
                    "bust": bust,
                    "waist": waist,
                    "hip": hip,
                    "occasion": occasion
                }
                user_states[user_id]["is_ready_for_outfit"] = True
                
                line_bot_api.reply_message(
                    reply_token,
                    TextSendMessage(text=f"已收到您的資訊！接下來請上傳三件上衣和三件褲子的圖片。請先輸入「上衣」或「褲子」來開始。")
                )
                return
        except ValueError:
            logging.error(f"Invalid user info format: '{text}'")
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="請依照格式輸入：身高,三圍胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
            )
            return

    # 處理衣物上傳模式
    if user_states[user_id]["is_ready_for_outfit"]:
        if text == "上衣":
            user_states[user_id]["current_mode"] = "upper"
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text=f"請上傳三件上衣圖片，您已上傳 {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE} 張。")
            )
            return
        elif text in ["褲子", "裙子"]:
            user_states[user_id]["current_mode"] = "lower"
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text=f"請上傳三件褲子/裙子圖片，您已上傳 {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE} 張。")
            )
            return
        
    # 如果沒有進入任何模式，給予提示
    line_bot_api.reply_message(
        reply_token,
        TextSendMessage(text="請先依照格式輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
    )

@line_handler.add(MessageEvent, message=ImageMessage)
def handle_image_message(event):
    user_id = event.source.user_id
    reply_token = event.reply_token
    
    logging.info(f"Received image message from user {user_id}")

    # 檢查是否已準備好處理圖片
    if not user_states[user_id]["is_ready_for_outfit"]:
        logging.warning("User is not in outfit collection mode.")
        line_bot_api.reply_message(
            reply_token,
            TextSendMessage(text="請先輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
        )
        return

    # 處理個人照片上傳
    if len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE and not user_states[user_id]["personal_photo"]:
        try:
            image_id = event.message.id
            base64_img = get_base64_image(image_id)
            
            # 先將個人照片的 Base64 儲存起來，用於 Gemini
            user_states[user_id]["personal_photo_base64"] = base64_img
            
            # 再將個人照片上傳到 Imgbb，並儲存 URL
            photo_url = upload_to_imgbb(base64_img)
            user_states[user_id]["personal_photo"] = photo_url
            user_states[user_id]["is_ready_for_photo"] = True
            
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="已收到您的個人照片，正在為您準備穿搭... 請稍候。")
            )
            
            # 準備 Gemini 提示
            user_info = user_states[user_id]["user_info"]
            occasion = user_info["occasion"]
            
            # Gemini 提示詞現在使用圖片 URL，而不是 Base64 編碼
            prompt = (
                f"我提供了三件上衣圖片和三件下半身圖片，以及使用者的一張個人照片。使用者資訊：身高 {user_info['height']}cm，三圍 {user_info['bust']}-{user_info['waist']}-{user_info['hip']}，場合是「{occasion}」。"
                "請根據這些衣物，為我推薦一套最適合的穿搭，並詳細說明為何這套搭配適合這個場合。請以繁體中文回答。"
                "請在描述中以圖片連結的形式（例如：'https://example.com/outfit1.jpg'）顯示你搭配好的衣物組合。"
            )
            
            all_images_base64 = (
                [get_base64_from_url(url) for url in user_states[user_id]["upper_body_images"]] +
                [get_base64_from_url(url) for url in user_states[user_id]["lower_body_images"]] +
                [user_states[user_id]["personal_photo_base64"]]
            )
            
            response_text = get_gemini_response(prompt, all_images_base64, GEMINI_API_URL)
            
            # 從 Gemini 建議中找出最佳搭配的衣物圖片（這裡需要更複雜的邏輯，我們使用佔位符）
            # 在實際應用中，您可以設計一個提示，讓 Gemini 返回最佳搭配的圖片索引
            best_upper_index = 0
            best_lower_index = 0

            best_upper_body_image = user_states[user_id]["upper_body_images"][best_upper_index]
            best_lower_body_image = user_states[user_id]["lower_body_images"][best_lower_index]

            # 呼叫虛擬試穿 API，並取得結果圖片 URL
            virtual_tryon_url = get_virtual_tryon_image(user_states[user_id]["personal_photo_base64"], get_base64_from_url(best_upper_body_image), get_base64_from_url(best_lower_body_image))
            
            # 發送 Gemini 的文字建議
            line_bot_api.push_message(
                user_id,
                TextSendMessage(text=f"這是為您推薦的搭配：\n\n{response_text}")
            )
            
            # 發送虛擬試穿照片
            if virtual_tryon_url:
                line_bot_api.push_message(
                    user_id,
                    ImageSendMessage(original_content_url=virtual_tryon_url, preview_image_url=virtual_tryon_url)
                )

            # 重置狀態以便下一次使用
            user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None, "is_ready_for_outfit": False, "is_ready_for_photo": False, "user_info": {}, "personal_photo": None, "personal_photo_base64": None})[user_id]
            return
            
        except Exception as e:
            logging.error(f"Error processing personal photo: {e}")
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text=f"圖片處理失敗，請稍後再試。錯誤：{e}")
            )
            return

    # 處理衣物圖片上傳
    mode = user_states[user_id]["current_mode"]
    if not mode:
        logging.warning("User is not in a defined upload mode.")
        line_bot_api.reply_message(
            reply_token,
            TextSendMessage(text="請先傳送「上衣」或「褲子」來選擇要上傳的衣服類型。")
        )
        return

    try:
        image_id = event.message.id
        base64_img = get_base64_image(image_id)
        # 上傳圖片到 Imgbb，並儲存 URL
        image_url = upload_to_imgbb(base64_img)
        if not image_url:
            logging.error("Failed to upload image to Imgbb.")
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="圖片上傳失敗，請稍後再試。")
            )
            return
            
        if mode == "upper":
            if len(user_states[user_id]["upper_body_images"]) < MAX_IMAGES_PER_TYPE:
                user_states[user_id]["upper_body_images"].append(image_url)
            else:
                logging.info("Upper body image count is full.")
                line_bot_api.reply_message(
                    reply_token,
                    TextSendMessage(text="上衣數量已滿，請傳送「褲子」來上傳褲子圖片。")
                )
                return
        else: # mode == "lower"
            if len(user_states[user_id]["lower_body_images"]) < MAX_IMAGES_PER_TYPE:
                user_states[user_id]["lower_body_images"].append(image_url)
            else:
                logging.info("Lower body image count is full.")
                line_bot_api.reply_message(
                    reply_token,
                    TextSendMessage(text="褲子數量已滿，請傳送「上衣」來上傳上衣圖片。")
                )
                return

        upper_count = len(user_states[user_id]["upper_body_images"])
        lower_count = len(user_states[user_id]["lower_body_images"])
        
        if upper_count < MAX_IMAGES_PER_TYPE or lower_count < MAX_IMAGES_PER_TYPE:
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text=f"已接收。上衣: {upper_count}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {lower_count}/{MAX_IMAGES_PER_TYPE}。")
            )
        else:
            logging.info("All outfit images received. Prompting for personal photo.")
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="已收到所有衣物圖片！接下來，請上傳一張您個人的全身照片，以便進行虛擬試穿。")
            )
            
    except Exception as e:
        logging.error(f"Error processing outfit image: {e}")
        line_bot_api.reply_message(
            reply_token,
            TextSendMessage(text=f"圖片處理失敗，請稍後再試。錯誤：{e}")
        )

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))