main.py

import os
import io
import re
from collections import defaultdict
import PIL.Image
import uvicorn
import requests
from pydantic_settings import BaseSettings

from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from google import genai
from google.genai import types

from linebot import LineBotApi, WebhookHandler
from linebot.exceptions import InvalidSignatureError
from linebot.models import (
    MessageEvent,
    TextMessage,
    TextSendMessage,
    ImageSendMessage,
    ImageMessage,
)

# LangChain 相關匯入
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


# ==========================
#  環境變數與設定管理
# ==========================

class Settings(BaseSettings):
    """使用 Pydantic 管理環境變數"""
    google_api_key: str
    channel_access_token: str
    channel_secret: str
    base_url: str  # 應用程式的公開網址，例如 ngrok 或 Hugging Face Space 的 URL

    class Config:
        env_file = ".env"

# 載入設定
settings = Settings()

# ==========================
#  API 客戶端與工具函式初始化
# ==========================

# 設置 Google AI API 金鑰
genai.configure(api_key=settings.google_api_key)

# 設置 Line Bot API
line_bot_api = LineBotApi(settings.channel_access_token)
line_handler = WebhookHandler(settings.channel_secret)

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

# 確保靜態檔案目錄存在
os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

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

def get_image_from_line(message_id: str) -> str | None:
    """
    從 Line 訊息 ID 獲取圖片內容並儲存到暫存檔案。
    """
    try:
        message_content = line_bot_api.get_message_content(message_id)
        # 使用 /tmp 目錄儲存暫存檔案，適合多數雲端環境
        file_path = f"/tmp/{message_id}.png"
        with open(file_path, "wb") as f:
            for chunk in message_content.iter_content():
                f.write(chunk)
        print(f"✅ 圖片成功儲存到：{file_path}")
        return file_path
    except Exception as e:
        print(f"❌ 從 Line 取得圖片失敗：{e}")
        return None


# ==========================
#  LangChain 工具定義
# ==========================

@tool
def generate_and_upload_image(prompt: str) -> str:
    """
    這個工具可以根據文字提示生成圖片，並將其上傳到伺服器。
    
    Args:
        prompt: 用於生成圖片的文字提示。
        
    Returns:
        回傳生成圖片的公開 URL。
    """
    try:
        # 使用 gemini-2.5-flash-image-preview 模型進行圖片生成
        model = genai.GenerativeModel('gemini-2.5-flash-image-preview')
        response = model.generate_content(
            contents=prompt,
            generation_config=types.GenerationConfig(response_modalities=['IMAGE'])
        )
        
        image_binary = None
        for part in response.candidates[0].content.parts:
            if part.inline_data and part.inline_data.data:
                image_binary = part.inline_data.data
                break
        
        if image_binary:
            image = PIL.Image.open(io.BytesIO(image_binary))
            # 隨機生成一個檔案名以避免衝突
            file_name = f"{os.urandom(16).hex()}.png"
            file_path = os.path.join("static", file_name)
            image.save(file_path, format="PNG")
            
            # 使用環境變數中的 BASE_URL 來建立完整的公開網址
            image_url = f"{settings.base_url}/{file_path}"
            print(f"✅ 圖片生成成功，URL: {image_url}")
            return f"圖片生成成功，請查看此 URL: {image_url}"
        
        return "圖片生成失敗，未收到有效的圖片資料。"
    except Exception as e:
        return f"圖片生成與上傳過程中發生錯誤: {e}"

@tool
def analyze_image_with_text(image_path: str, user_text: str) -> str:
    """
    這個工具可以根據圖片和文字提示來回答問題。
    
    Args:
        image_path: 圖片在本地端儲存的路徑。
        user_text: 針對圖片提出的文字問題。
        
    Returns:
        模型針對圖片和文字提示給出的回應。
    """
    try:
        if not os.path.exists(image_path):
            return "圖片路徑無效，無法進行分析。"
            
        img_user = PIL.Image.open(image_path)
        model = genai.GenerativeModel("gemini-1.5-flash") # 使用 gemini 1.5 flash 模型
        response = model.generate_content([img_user, user_text])

        if response.text:
            return response.text
        else:
            return "Gemini 沒有給出答案，請嘗試換個方式提問！"
    except Exception as e:
        return f"圖片分析過程中發生錯誤: {e}"


# ==========================
#  LangChain 代理人設定
# ==========================

# 結合所有工具
tools = [generate_and_upload_image, analyze_image_with_text]

# 建立 LLM 模型實例
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.3)

# 建立提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一個強大的圖像生成與問答助理。
- 當用戶的指令明顯是要生成圖片時 (例如：'畫一張...'、'生成...'、'幫我做一張圖...')，請使用 `generate_and_upload_image` 工具。
- 當用戶的指令包含圖片路徑 (image_path) 和問題時，請使用 `analyze_image_with_text` 工具。
- 成功執行 `generate_and_upload_image` 工具後，你會獲得一個 URL，你的最終回答必須包含這個 URL。
- 如果工具執行過程中產生任何錯誤訊息，請以友善的方式解讀並回應給用戶。"""),
    ("user", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])


# 建立代理人
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# ==========================
#  FastAPI 路由
# ==========================

@app.get("/")
def root():
    return {"message": "Line Bot is running!"}

@app.post("/webhook")
async def webhook(
    request: Request,
    background_tasks: BackgroundTasks,
    x_line_signature: str = Header(None),
):
    body = await request.body()
    try:
        # 使用背景任務處理 Webhook，避免 Line Server 超時
        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 isinstance(event.message, ImageMessage):
        image_path = get_image_from_line(event.message.id)
        if image_path:
            try:
                # 組合給代理人的輸入，使用一個通用的問題來分析圖片
                agent_input = {
                    "input": f"這是一張使用者上傳的圖片，請詳細描述你看到了什麼。圖片的路徑是 '{image_path}'。"
                }
                # 運行代理人
                response = agent_executor.invoke(agent_input)
                output_text = response["output"]
                
                # 回覆分析結果
                line_bot_api.reply_message(event.reply_token, TextSendMessage(text=output_text))

            except Exception as e:
                print(f"代理人執行出錯: {e}")
                error_message = f"圖片分析時發生錯誤，請稍後再試。\n錯誤訊息：{e}"
                line_bot_api.reply_message(event.reply_token, TextSendMessage(text=error_message))
        else:
            line_bot_api.reply_message(
                event.reply_token, TextSendMessage(text="❌ 圖片接收失敗，請再試一次。")
            )

    # 處理文字訊息：主要用於生成圖片或一般問答
    elif isinstance(event.message, TextMessage):
        user_text = event.message.text
        agent_input = {"input": user_text}
        
        try:
            # 運行代理人
            response = agent_executor.invoke(agent_input)
            output_text = response["output"]

            # 使用正規表示法尋找 URL，更穩定
            image_url_match = re.search(r'https?://\S+\.(?:png|jpg|jpeg|gif)', output_text, re.IGNORECASE)
            
            if image_url_match:
                image_url = image_url_match.group(0)
                # 推送訊息，包含生成的圖片
                line_bot_api.push_message(
                    event.source.user_id,
                    [
                        TextSendMessage(text="✨ 這是我為您生成的圖片喔～"),
                        ImageSendMessage(original_content_url=image_url, preview_image_url=image_url)
                    ]
                )
            else:
                # 若無圖片 URL，則直接回覆文字
                line_bot_api.reply_message(event.reply_token, TextSendMessage(text=output_text))
        except Exception as e:
            print(f"代理人執行出錯: {e}")
            error_message = f"代理人執行時發生錯誤，請稍後再試。\n錯誤訊息：{e}"
            line_bot_api.reply_message(event.reply_token, TextSendMessage(text=error_message))

if __name__ == "__main__":
    # 使用 settings 物件中的設定
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)