main.py

from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from linebot import LineBotApi, WebhookHandler
from linebot.exceptions import InvalidSignatureError
from linebot.models import (
    MessageEvent,
    TextMessage,
    TextSendMessage,
    ImageSendMessage,
    ImageMessage,
)
from google import genai
from google.genai import types
from PIL import Image
from collections import defaultdict
import os
import io
import requests
import uvicorn
import logging
import base64
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import AgentExecutor, create_tool_calling_agent
from pydantic import BaseModel, Field
from typing import List

# ==========================#  環境設定與工具函式# ==========================#
# 設置日誌記錄，級別為 INFO，格式包含時間、級別和訊息
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", "HF_SPACE"]):
    logging.error("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and HF_SPACE.")
    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and HF_SPACE.")

# 獲取環境變數中的金鑰和 URL
google_api = os.environ["GOOGLE_API_KEY"]
line_channel_access_token = os.environ["CHANNEL_ACCESS_TOKEN"]
line_channel_secret = os.environ["CHANNEL_SECRET"]
HF_SPACE_URL = os.environ["HF_SPACE"]

# Line Bot API 設定
line_bot_api = LineBotApi(line_channel_access_token)
line_handler = WebhookHandler(line_channel_secret)

# Google AI API 設定
genai_client = genai.Client(api_key=google_api)

# 使用者狀態追蹤，用來儲存已上傳的衣物圖片 URL
user_states = defaultdict(lambda: {
    "upper_body_images": [],
    "lower_body_images": [],
    "current_mode": None,
})

MAX_IMAGES_PER_TYPE = 3
IMAGIN_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent"

# 建立 FastAPI 應用程式
app = FastAPI()

# 設定靜態文件服務，用於託管圖片
app.mount("/static", StaticFiles(directory="static"), name="static")

# 設定 CORS 跨域請求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

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 save_image_locally(image_binary: bytes):
    """
    將二進位圖片資料儲存到本地，並返回一個可供外部存取的 URL。
    """
    try:
        # 確保 'static' 資料夾存在
        if not os.path.exists("static"):
            os.makedirs("static")
        
        image = Image.open(io.BytesIO(image_binary))
        # 隨機生成一個檔案名以避免衝突
        file_name = f"static/{os.urandom(16).hex()}.png"
        image.save(file_name, format="PNG")
        
        image_url = os.path.join(HF_SPACE_URL, file_name)
        logging.info(f"Image successfully saved locally: {image_url}")
        return image_url
    except Exception as e:
        logging.error(f"Error saving image locally: {e}")
        return None

def get_image_url_from_line(message_id):
    """
    從 Line 訊息 ID 獲取圖片內容並儲存到暫存檔案。
    """
    try:
        message_content = line_bot_api.get_message_content(message_id)
        # 取得二進位圖片資料
        image_binary = message_content.content
        # 使用 save_image_locally 函式儲存圖片並取得 URL
        return save_image_locally(image_binary)
    except Exception as e:
        print(f"❌ 圖片取得失敗：{e}")
        return None

# ==========================#  LangChain 工具定義# ==========================#
# 定義工具的輸入模型，明確指定參數的型別和描述
class OutfitInput(BaseModel):
    """用於生成穿搭圖片的輸入參數。"""
    upper_body_urls: List[str] = Field(..., description="上衣圖片的 URL 列表。")
    lower_body_urls: List[str] = Field(..., description="褲子/裙子圖片的 URL 列表。")

@tool(args_schema=OutfitInput)
def generate_outfit_from_clothes(upper_body_urls: list, lower_body_urls: list) -> str:
    """
    這個工具可以根據提供的上衣和褲子/裙子圖片 URLs，生成一套全新的穿搭圖片。
    
    Args:
        upper_body_urls: 一組上衣圖片的 URL 列表。
        lower_body_urls: 一組褲子/裙子圖片的 URL 列表。
        
    Returns:
        回傳生成圖片的 URL。
    """
    logging.info("Attempting to generate a new outfit image.")
    prompt = "使用提供的上衣和褲子/裙子圖片，生成一套完整且時尚的穿搭圖片。請將衣服呈現在一個有模特兒穿著或是在平面上呈現的完整畫面中。風格應與提供的衣物相符。"
    
    all_images_base64 = (
        [get_base64_from_url(url) for url in upper_body_urls] +
        [get_base64_from_url(url) for url in lower_body_urls]
    )

    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={os.environ['GOOGLE_API_KEY']}", json=payload, timeout=120)
        response.raise_for_status()
        response_data = response.json()
        if 'candidates' in response_data and response_data['candidates']:
            image_part = response_data['candidates'][0]['content']['parts'][0]
            if image_part and 'inlineData' in image_part:
                generated_image_base64 = image_part['inlineData']['data']
                generated_image_url = save_image_locally(base64.b64decode(generated_image_base64))
                if generated_image_url:
                    logging.info(f"Successfully generated and saved new outfit image: {generated_image_url}")
                    return generated_image_url
                else:
                    logging.error("Failed to save generated outfit image locally.")
                    return None
            else:
                logging.error("Gemini image response format is invalid for outfit generation.")
                return None
        else:
            logging.error(f"Gemini outfit generation response has no candidates: {response.text}")
            return None
    except requests.exceptions.RequestException as e:
        logging.error(f"Gemini outfit generation API request failed: {e}")
        return None

# ==========================#  LangChain 代理人設定# ==========================#
# 結合所有工具
tools = [generate_outfit_from_clothes]
# 建立 LLM 模型實例
llm = ChatGoogleGenerativeAI(google_api_key=google_api, model="gemini-2.5-flash", temperature=0.2)
# 建立提示模板
prompt_template = ChatPromptTemplate([
    ("system", "你是一個強大的圖像生成與問答助理，可以根據用戶的請求使用提供的工具。當你執行 generate_outfit_from_clothes 工具成功後會獲得一個 URL，然後你回答的 output 要包含有這個 URL 的完整資訊。如果工具有產生錯誤訊息請解讀並回應。"),
    ("user", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])
# 建立代理人
agent = create_tool_calling_agent(llm, tools, prompt_template)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# ==========================#  FastAPI 路由# ==========================#
@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"

@line_handler.add(MessageEvent, message=(ImageMessage, TextMessage))
def handle_message(event):
    user_id = event.source.user_id
    
    # 處理圖片上傳
    if event.message.type == "image":
        image_url = get_image_url_from_line(event.message.id)
        if image_url:
            mode = user_states[user_id]["current_mode"]
            if not mode:
                line_bot_api.reply_message(
                    event.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)
                    reply_text = f"已接收上衣。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text=reply_text))
                else:
                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text="上衣數量已滿。"))
            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)
                    reply_text = f"已接收褲子/裙子。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text=reply_text))
                else:
                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text="褲子/裙子數量已滿。"))

            # 檢查是否所有圖片都已上傳完畢
            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):
                line_bot_api.push_message(
                    user_id,
                    TextSendMessage(text="所有衣物圖片已收集完畢！\n\n現在您可以輸入「**生成穿搭**」或「**圖片推薦**」來獲得一套新的圖片穿搭。")
                )
        else:
            line_bot_api.reply_message(
                event.reply_token, TextSendMessage(text="沒有接收到圖片~")
            )
            
    # 處理文字訊息
    elif event.message.type == "text":
        user_text = event.message.text.lower()
        reply_token = event.reply_token
        
        # 處理重置功能
        if user_text in ["重置", "重來", "重新開始", "再一次"]:
            user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None})[user_id]
            line_bot_api.reply_message(
                reply_token, TextSendMessage(text="狀態已重置。請先輸入「上衣」或「褲子」來開始。")
            )
            return
            
        # 處理衣物上傳模式切換
        if user_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 user_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

        # 處理圖片生成穿搭
        if user_text in ["生成穿搭", "圖片推薦"]:
            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):
                try:
                    line_bot_api.reply_message(reply_token, TextSendMessage(text="好的，正在為您生成一套新的穿搭圖片，請稍候..."))
                    
                    # 直接呼叫生成工具，並傳入已收集的圖片 URLs
                    generated_image_url = generate_outfit_from_clothes(
                        user_states[user_id]["upper_body_images"],
                        user_states[user_id]["lower_body_images"]
                    )
                    
                    if generated_image_url:
                        line_bot_api.push_message(
                            user_id,
                            ImageSendMessage(original_content_url=generated_image_url, preview_image_url=generated_image_url)
                        )
                        line_bot_api.push_message(
                            user_id,
                            TextSendMessage(text="這是根據您的衣物生成的圖片推薦。如果想再次使用，請輸入「重置」。")
                        )
                    else:
                         line_bot_api.push_message(user_id, TextSendMessage(text="圖片生成失敗，請稍後再試。"))
                    
                except Exception as e:
                    logging.error(f"Error generating outfit image: {e}")
                    line_bot_api.push_message(user_id, TextSendMessage(text=f"圖片生成失敗，請稍後再試。錯誤：{e}"))

                # 生成後重置狀態以便下一次使用
                user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None})[user_id]
                return
            else:
                line_bot_api.reply_message(
                    reply_token, TextSendMessage(text="請先上傳三件上衣和三件褲子/裙子圖片，再輸入「生成穿搭」來獲得圖片推薦。")
                )
                return

        # 如果都不是特定指令，則交給代理人處理
        agent_input = {"input": user_text}
        try:
            # 運行代理人
            response = agent_executor.invoke(agent_input)
            out = response["output"]
            line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
        except Exception as e:
            print(f"代理人執行出錯: {e}")
            out = f"代理人執行出錯!錯誤訊息：{e}"
            line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)