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

# 檢查環境變數
if not all(k in os.environ for k in ["CHANNEL_ACCESS_TOKEN", "CHANNEL_SECRET", "GOOGLE_API_KEY"]):
    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, and GOOGLE_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 # 儲存個人照片
})

MAX_IMAGES_PER_TYPE = 3
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-05-20:generateContent"
GOOGLE_API_KEY = os.environ["GOOGLE_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):
    message_content = line_bot_api.get_message_content(message_id)
    image_bytes = message_content.content
    return base64.b64encode(image_bytes).decode('utf-8')

def get_gemini_response(prompt: str, images: list):
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ] + [
                    {"inlineData": {"mimeType": "image/jpeg", "data": img}} for img in images
                ]
            }
        ]
    }
    
    response = requests.post(f"{GEMINI_API_URL}?key={GOOGLE_API_KEY}", json=payload)
    if response.status_code == 200:
        return response.json()['candidates'][0]['content']['parts'][0]['text']
    else:
        return f"Gemini API 請求失敗：{response.status_code}, {response.text}"

@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
    
    # 處理個人資訊輸入
    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:
            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 == "褲子":
            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
        elif text == "重置":
            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})[user_id]
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="狀態已重置。請重新輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合，例如：165,85,65,90,約會")
            )
            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
    
    # 檢查是否已準備好處理圖片
    if not user_states[user_id]["is_ready_for_outfit"]:
        line_bot_api.reply_message(
            reply_token,
            TextSendMessage(text="請先輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
        )
        return

    # 如果個人照片還沒上傳
    if not user_states[user_id]["personal_photo"]:
        try:
            image_id = event.message.id
            base64_img = get_base64_image(image_id)
            user_states[user_id]["personal_photo"] = base64_img
            user_states[user_id]["is_ready_for_photo"] = True
            
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="已收到您的個人照片，正在為您準備穿搭... 請稍候。")
            )
            
            # 開始生成搭配建議和試穿照片
            user_info = user_states[user_id]["user_info"]
            occasion = user_info["occasion"]
            
            prompt = (
                f"我提供了三件上衣圖片和三件下半身圖片，以及使用者的一張個人照片。使用者的身高為 {user_info['height']}，三圍為 {user_info['bust']}-{user_info['waist']}-{user_info['hip']}，想要參加的場合是「{occasion}」。"
                "請根據這些衣物，為我推薦一套最適合的穿搭，並詳細說明為何這套搭配適合這個場合。請以繁體中文回答。"
                "接著，請生成一張虛擬試穿的照片，將推薦的上衣和下衣搭配到使用者提供的照片上。由於模型限制，我會使用佔位符圖片來模擬試穿效果。"
            )
            
            all_images = user_states[user_id]["upper_body_images"] + user_states[user_id]["lower_body_images"] + [user_states[user_id]["personal_photo"]]
            
            response_text = get_gemini_response(prompt, all_images)
            
            # 虛擬試穿照片佔位符
            virtual_try_on_url = "https://placehold.co/1024x1024?text=Virtual+Try-On+Outfit"
            
            # 發送 Gemini 的文字建議
            line_bot_api.push_message(
                user_id,
                TextSendMessage(text=f"這是為您推薦的搭配：\n\n{response_text}")
            )
            
            # 發送虛擬試穿照片
            line_bot_api.push_message(
                user_id,
                ImageSendMessage(original_content_url=virtual_try_on_url, preview_image_url=virtual_try_on_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})[user_id]
            return
            
        except Exception as e:
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text=f"圖片處理失敗，請稍後再試。錯誤：{e}")
            )
            return

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

    try:
        image_id = event.message.id
        base64_img = get_base64_image(image_id)
        
        if mode == "upper":
            if len(user_states[user_id]["upper_body_images"]) < MAX_IMAGES_PER_TYPE:
                user_states[user_id]["upper_body_images"].append(base64_img)
            else:
                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(base64_img)
            else:
                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:
            line_bot_api.reply_message(
                reply_token,
                TextSendMessage(text="已收到所有衣物圖片！接下來，請上傳一張您個人的全身照片，以便進行虛擬試穿。如果想重來請輸入：重置")
            )
            
    except Exception as 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)))