main.py

from fastapi import Request, FastAPI, HTTPException
import os
import sys
import asyncio
import aiohttp
import aiofiles
from pathlib import Path
from typing import Optional

from linebot.models import (
    MessageEvent, TextSendMessage, FileMessage, ImageMessage,
    PostbackEvent, TemplateSendMessage, CarouselTemplate, CarouselColumn,
    PostbackAction, QuickReply, QuickReplyButton, MessageAction
)
from linebot.exceptions import InvalidSignatureError
from linebot.aiohttp_async_http_client import AiohttpAsyncHttpClient
from linebot import AsyncLineBotApi, WebhookParser

# Google GenAI imports
from google import genai
from google.genai import types

# File Manager Agent
from file_manager_agent import FileManagerAgent

# Configuration
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") or ""

# LINE Bot configuration
channel_secret = os.getenv("ChannelSecret", None)
channel_access_token = os.getenv("ChannelAccessToken", None)

# Validate environment variables
if channel_secret is None:
    print("Specify ChannelSecret as environment variable.")
    sys.exit(1)
if channel_access_token is None:
    print("Specify ChannelAccessToken as environment variable.")
    sys.exit(1)
if not GOOGLE_API_KEY:
    raise ValueError("Please set GOOGLE_API_KEY via env var or code.")

# Initialize GenAI client (Note: File Search API only supports Gemini API, not VertexAI)
client = genai.Client(api_key=GOOGLE_API_KEY)

print("GenAI client initialized successfully.")

# Initialize the FastAPI app for LINEBot
app = FastAPI()
client_session = aiohttp.ClientSession()
async_http_client = AiohttpAsyncHttpClient(client_session)
line_bot_api = AsyncLineBotApi(channel_access_token, async_http_client)
parser = WebhookParser(channel_secret)

# Create uploads directory if not exists
UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)

# Model configuration
MODEL_NAME = "gemini-2.5-flash"

def get_store_name(event: MessageEvent) -> str:
    """
    Get the file search store name based on the message source.
    Returns user_id for 1-on-1 chat, group_id for group chat.
    """
    if event.source.type == "user":
        return f"user_{event.source.user_id}"
    elif event.source.type == "group":
        return f"group_{event.source.group_id}"
    elif event.source.type == "room":
        return f"room_{event.source.room_id}"
    else:
        return f"unknown_{event.source.user_id}"


async def download_line_content(message_id: str, file_name: str) -> Optional[Path]:
    """
    Download file content from LINE and save to local uploads directory.
    Returns the local file path if successful, None otherwise.
    """
    try:
        # Get message content from LINE
        message_content = await line_bot_api.get_message_content(message_id)

        # Extract file extension from original file name
        _, ext = os.path.splitext(file_name)
        # Use safe file name (ASCII only) to avoid encoding issues
        safe_file_name = f"{message_id}{ext}"
        file_path = UPLOAD_DIR / safe_file_name

        async with aiofiles.open(file_path, 'wb') as f:
            async for chunk in message_content.iter_content():
                await f.write(chunk)

        print(f"Downloaded file: {file_path} (original: {file_name})")
        return file_path
    except Exception as e:
        print(f"Error downloading file: {e}")
        return None


async def ensure_file_search_store_exists(store_name: str) -> tuple[bool, str]:
    """
    Ensure file search store exists, create if not.
    Returns (success, actual_store_name).
    Note: store_name is used as display_name, but actual name is auto-generated by API.
    """
    try:
        # List all stores and check if one with our display_name exists
        stores = client.file_search_stores.list()
        for store in stores:
            if hasattr(store, 'display_name') and store.display_name == store_name:
                print(f"File search store '{store_name}' already exists: {store.name}")
                return True, store.name

        # Store doesn't exist, create it
        print(f"Creating file search store with display_name '{store_name}'...")
        store = client.file_search_stores.create(
            config={'display_name': store_name}
        )
        print(f"File search store created: {store.name} (display_name: {store_name})")
        return True, store.name

    except Exception as e:
        print(f"Error ensuring file search store exists: {e}")
        return False, ""


# Cache to store display_name -> actual_name mapping
store_name_cache = {}

# Cache to store citations/grounding metadata for each user/group
# Key: store_name, Value: list of grounding chunks
citations_cache = {}


async def list_documents_in_store(store_name: str) -> list:
    """
    List all documents in a file search store.
    Returns list of document info dicts.
    """
    try:
        # Get actual store name
        actual_store_name = None
        if store_name in store_name_cache:
            actual_store_name = store_name_cache[store_name]
        else:
            # Find store by display_name
            stores = client.file_search_stores.list()
            for store in stores:
                if hasattr(store, 'display_name') and store.display_name == store_name:
                    actual_store_name = store.name
                    store_name_cache[store_name] = actual_store_name
                    break

        if not actual_store_name:
            print(f"Store '{store_name}' not found")
            return []

        documents = []

        # Try to use SDK method first
        if hasattr(client.file_search_stores, 'documents'):
            for doc in client.file_search_stores.documents.list(parent=actual_store_name):
                documents.append({
                    'name': doc.name,
                    'display_name': getattr(doc, 'display_name', 'Unknown'),
                    'create_time': str(getattr(doc, 'create_time', '')),
                    'update_time': str(getattr(doc, 'update_time', ''))
                })
                print(f"Use SDK list function: File found in store '{store_name}': {doc.name}")
        else:
            # Fallback to REST API
            import requests
            url = f"https://generativelanguage.googleapis.com/v1beta/{actual_store_name}/documents"
            headers = {'Content-Type': 'application/json'}
            params = {'key': GOOGLE_API_KEY}

            response = requests.get(url, headers=headers, params=params, timeout=10)
            response.raise_for_status()
            data = response.json()

            for doc in data.get('documents', []):
                documents.append({
                    'name': doc.get('name', 'N/A'),
                    'display_name': doc.get('displayName', 'Unknown'),
                    'create_time': doc.get('createTime', ''),
                    'update_time': doc.get('updateTime', '')
                })
                print(f"Use REST API list function: File found in store '{store_name}': {doc.name}")
        return documents

    except Exception as e:
        print(f"Error listing documents in store: {e}")
        return []


async def delete_document(document_name: str) -> bool:
    """
    Delete a document from file search store.
    Returns True if successful, False otherwise.
    Note: force=True is required to permanently delete documents from File Search Store.
    """
    try:
        # Try to use SDK method first with force=True
        try:
            if hasattr(client.file_search_stores, 'documents'):
                # Force delete is required for File Search Store documents
                client.file_search_stores.documents.delete(
                    name=document_name,
                    config={'force': True}
                )
                print(f"Document deleted successfully with force=True: {document_name}")
                return True
        except Exception as sdk_error:
            print(f"SDK delete failed, trying REST API: {sdk_error}")

        # Fallback to REST API with force parameter
        import requests
        url = f"https://generativelanguage.googleapis.com/v1beta/{document_name}"
        headers = {'Content-Type': 'application/json'}
        params = {
            'key': GOOGLE_API_KEY,
            'force': 'true'  # Required for File Search Store documents
        }

        response = requests.delete(url, headers=headers, params=params, timeout=10)
        response.raise_for_status()

        print(f"Document deleted successfully via REST API with force=true: {document_name}")
        return True

    except Exception as e:
        print(f"Error deleting document: {e}")
        return False


async def upload_to_file_search_store(file_path: Path, store_name: str, display_name: Optional[str] = None) -> bool:
    """
    Upload a file to Gemini file search store.
    Returns True if successful, False otherwise.
    """
    try:
        # Check cache first
        if store_name in store_name_cache:
            actual_store_name = store_name_cache[store_name]
            print(f"Using cached store name: {actual_store_name}")
        else:
            # Ensure the store exists before uploading
            success, actual_store_name = await ensure_file_search_store_exists(store_name)
            if not success:
                print(f"Failed to ensure store '{store_name}' exists")
                return False
            # Cache the mapping
            store_name_cache[store_name] = actual_store_name

        # Upload to file search store
        # actual_store_name is the API-generated name (e.g., fileSearchStores/xxx)
        # display_name is the custom display name for the file (used in citations)
        config_dict = {}
        if display_name:
            config_dict['display_name'] = display_name

        operation = client.file_search_stores.upload_to_file_search_store(
            file_search_store_name=actual_store_name,
            file=str(file_path),
            config=config_dict if config_dict else None
        )

        # Wait for operation to complete (with timeout)
        max_wait = 60  # seconds
        elapsed = 0
        while not operation.done and elapsed < max_wait:
            await asyncio.sleep(2)
            operation = client.operations.get(operation)
            elapsed += 2

        if operation.done:
            print(f"File uploaded to store '{store_name}': {operation}")
            return True
        else:
            print(f"Upload operation timeout for store '{store_name}'")
            return False

    except Exception as e:
        print(f"Error uploading to file search store: {e}")
        return False


async def query_file_search(query: str, store_name: str) -> tuple[str, list]:
    """
    Query the file search store using generate_content.
    Returns (AI response text, list of citations).
    """
    try:
        # Get actual store name from cache or by searching
        actual_store_name = None

        if store_name in store_name_cache:
            actual_store_name = store_name_cache[store_name]
            print(f"Using cached store name for query: {actual_store_name}")
        else:
            # Try to find the store by display_name
            try:
                stores = client.file_search_stores.list()
                for store in stores:
                    if hasattr(store, 'display_name') and store.display_name == store_name:
                        actual_store_name = store.name
                        store_name_cache[store_name] = actual_store_name
                        print(f"Found store for query: {actual_store_name}")
                        break
            except Exception as list_error:
                print(f"Error listing stores: {list_error}")

        if not actual_store_name:
            # Store doesn't exist - guide user to upload files
            print(f"File search store '{store_name}' not found")
            return ("📁 您還沒有上傳任何檔案。\n\n請先傳送文件檔案（PDF、DOCX、TXT 等）給我，上傳完成後就可以開始提問了！\n\n💡 提示：如果您想分析圖片，請直接傳送圖片給我，我會立即為您分析。", [])

        # Create FileSearch tool with actual store name
        tool = types.Tool(
            file_search=types.FileSearch(
                file_search_store_names=[actual_store_name]
            )
        )

        # Generate content with file search
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=query,
            config=types.GenerateContentConfig(
                tools=[tool],
                temperature=0.7,
            )
        )

        # Extract grounding metadata (citations)
        citations = []
        try:
            if hasattr(response, 'candidates') and response.candidates:
                candidate = response.candidates[0]
                if hasattr(candidate, 'grounding_metadata') and candidate.grounding_metadata:
                    grounding_chunks = candidate.grounding_metadata.grounding_chunks
                    for chunk in grounding_chunks:
                        if hasattr(chunk, 'web') and chunk.web:
                            # Web source
                            citations.append({
                                'type': 'web',
                                'title': getattr(chunk.web, 'title', 'Unknown'),
                                'uri': getattr(chunk.web, 'uri', ''),
                            })
                        elif hasattr(chunk, 'retrieved_context') and chunk.retrieved_context:
                            # File search source
                            citations.append({
                                'type': 'file',
                                'title': getattr(chunk.retrieved_context, 'title', 'Unknown'),
                                'text': getattr(chunk.retrieved_context, 'text', '')[:500],  # Limit to 500 chars
                            })
            print(f"Found {len(citations)} citations")
        except Exception as citation_error:
            print(f"Error extracting citations: {citation_error}")

        # Extract text from response
        if response.text:
            return (response.text, citations)
        else:
            return ("抱歉，我無法從文件中找到相關資訊。", [])

    except Exception as e:
        print(f"Error querying file search: {e}")
        # Check if error is related to missing store
        if "not found" in str(e).lower() or "does not exist" in str(e).lower():
            return ("📁 您還沒有上傳任何檔案。\n\n請先傳送文件檔案（PDF、DOCX、TXT 等）給我，上傳完成後就可以開始提問了！\n\n💡 提示：如果您想分析圖片，請直接傳送圖片給我，我會立即為您分析。", [])
        return (f"查詢時發生錯誤：{str(e)}", [])


async def analyze_image_with_gemini(image_path: Path) -> str:
    """
    Analyze image using Gemini's vision capability.
    Returns the analysis result text.
    """
    try:
        # Read image bytes
        with open(image_path, 'rb') as f:
            image_bytes = f.read()

        # Determine MIME type based on file extension
        ext = image_path.suffix.lower()
        mime_type_map = {
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.png': 'image/png',
            '.gif': 'image/gif',
            '.webp': 'image/webp'
        }
        mime_type = mime_type_map.get(ext, 'image/jpeg')

        # Create image part
        image = types.Part.from_bytes(
            data=image_bytes,
            mime_type=mime_type
        )

        # Generate content with image
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=["請詳細描述這張圖片的內容，包括主要物品、場景、文字等資訊。", image],
        )

        if response.text:
            return response.text
        else:
            return "抱歉，我無法分析這張圖片。"

    except Exception as e:
        print(f"Error analyzing image with Gemini: {e}")
        return f"圖片分析時發生錯誤：{str(e)}"


async def handle_image_message(event: MessageEvent, message: ImageMessage):
    """
    Handle image messages - analyze using Gemini vision.
    """
    file_name = f"image_{message.id}.jpg"

    # Download image
    reply_msg = TextSendMessage(text="正在分析您的圖片，請稍候...")
    await line_bot_api.reply_message(event.reply_token, reply_msg)

    file_path = await download_line_content(message.id, file_name)

    if file_path is None:
        error_msg = TextSendMessage(text="圖片下載失敗，請重試。")
        await line_bot_api.push_message(event.source.user_id, error_msg)
        return

    # Analyze image with Gemini
    analysis_result = await analyze_image_with_gemini(file_path)

    # Clean up local file
    try:
        file_path.unlink()
    except Exception as e:
        print(f"Error deleting file: {e}")

    # Send analysis result
    result_msg = TextSendMessage(text=f"📸 圖片分析結果：\n\n{analysis_result}")
    await line_bot_api.push_message(event.source.user_id, result_msg)


async def handle_document_message(event: MessageEvent, message: FileMessage):
    """
    Handle file messages - download and upload to file search store.
    """
    store_name = get_store_name(event)
    file_name = message.file_name or "unknown_file"

    # Download file
    reply_msg = TextSendMessage(text="正在處理您的檔案，請稍候...")
    await line_bot_api.reply_message(event.reply_token, reply_msg)

    file_path = await download_line_content(message.id, file_name)

    if file_path is None:
        error_msg = TextSendMessage(text="檔案下載失敗，請重試。")
        await line_bot_api.push_message(event.source.user_id, error_msg)
        return

    # Upload to file search store
    success = await upload_to_file_search_store(file_path, store_name, file_name)

    # Clean up local file
    try:
        file_path.unlink()
    except Exception as e:
        print(f"Error deleting file: {e}")

    if success:
        # Create Quick Reply buttons for common actions with specific file name
        quick_reply = QuickReply(items=[
            QuickReplyButton(action=MessageAction(label="📝 生成檔案摘要", text=f"請幫我生成「{file_name}」這個檔案的摘要")),
            QuickReplyButton(action=MessageAction(label="📌 重點整理", text=f"請幫我整理「{file_name}」的重點")),
            QuickReplyButton(action=MessageAction(label="📋 列出檔案", text="列出檔案")),
        ])

        success_msg = TextSendMessage(
            text=f"✅ 檔案已成功上傳！\n檔案名稱：{file_name}\n\n現在您可以詢問我關於這個檔案的任何問題。",
            quick_reply=quick_reply
        )
        await line_bot_api.push_message(event.source.user_id, success_msg)
    else:
        error_msg = TextSendMessage(text="檔案上傳失敗，請重試。")
        await line_bot_api.push_message(event.source.user_id, error_msg)


def is_list_files_intent(text: str) -> bool:
    """
    Check if user wants to list files.
    """
    list_keywords = [
        '列出檔案', '列出文件', '顯示檔案', '顯示文件',
        '查看檔案', '查看文件', '檔案列表', '文件列表',
        '有哪些檔案', '有哪些文件', '我的檔案', '我的文件',
        'list files', 'show files', 'my files'
    ]
    text_lower = text.lower().strip()
    return any(keyword in text_lower for keyword in list_keywords)


async def send_files_carousel(event: MessageEvent, documents: list):
    """
    Send files as LINE Carousel Template.
    """
    if not documents:
        no_files_msg = TextSendMessage(text="📁 目前沒有任何文件。\n\n請先上傳文件檔案，就可以查詢囉！")
        await line_bot_api.reply_message(event.reply_token, no_files_msg)
        return

    # LINE Carousel限制最多10個
    documents = documents[:10]

    columns = []
    for doc in documents:
        # 提取檔名（去除路徑部分）
        display_name = doc.get('display_name', 'Unknown')
        # 格式化時間
        create_time = doc.get('create_time', '')
        if create_time and 'T' in create_time:
            # 簡化時間顯示 (YYYY-MM-DD HH:MM)
            try:
                from datetime import datetime
                dt = datetime.fromisoformat(create_time.replace('Z', '+00:00'))
                create_time = dt.strftime('%Y-%m-%d %H:%M')
            except:
                create_time = create_time[:16]  # 簡單截斷

        # 建立每個檔案的 Column
        column = CarouselColumn(
            thumbnail_image_url='https://via.placeholder.com/1024x1024/4CAF50/FFFFFF?text=File',  # 預設圖片
            title=display_name[:40],  # LINE 限制標題長度
            text=f"上傳時間：{create_time[:20]}" if create_time else "文件檔案",
            actions=[
                PostbackAction(
                    label='🗑️ 刪除檔案',
                    data=f"action=delete_file&doc_name={doc['name']}"
                )
            ]
        )
        columns.append(column)

    carousel_template = CarouselTemplate(columns=columns)
    template_message = TemplateSendMessage(
        alt_text=f'📁 找到 {len(documents)} 個文件',
        template=carousel_template
    )

    await line_bot_api.reply_message(event.reply_token, template_message)


async def handle_postback(event: PostbackEvent):
    """
    Handle postback events (e.g., delete file button clicks).
    """
    try:
        # Parse postback data
        data = event.postback.data
        params = dict(param.split('=') for param in data.split('&'))

        action = params.get('action')
        doc_name = params.get('doc_name')

        if action == 'delete_file' and doc_name:
            # Delete the document
            success = await delete_document(doc_name)

            if success:
                # Extract display name from doc_name for user-friendly message
                display_name = doc_name.split('/')[-1] if '/' in doc_name else doc_name

                reply_msg = TextSendMessage(
                    text=f"✅ 檔案已刪除成功！\n\n如需查看剩餘檔案，請輸入「列出檔案」。"
                )
            else:
                reply_msg = TextSendMessage(text="❌ 刪除檔案失敗，請稍後再試。")

            await line_bot_api.reply_message(event.reply_token, reply_msg)
        else:
            print(f"Unknown postback action: {action}")

    except Exception as e:
        print(f"Error handling postback: {e}")
        error_msg = TextSendMessage(text="處理操作時發生錯誤。")
        await line_bot_api.reply_message(event.reply_token, error_msg)


async def handle_text_message(event: MessageEvent, message):
    """
    Handle text messages - query the file search store or list files.
    """
    store_name = get_store_name(event)
    query = message.text

    print(f"Received query: {query} for store: {store_name}")

    # Check if user wants to view a citation
    if query.startswith("📖 引用"):
        # Extract citation number
        try:
            citation_num = int(query.replace("📖 引用", "").strip())
            if store_name in citations_cache and 0 < citation_num <= len(citations_cache[store_name]):
                citation = citations_cache[store_name][citation_num - 1]

                # Format citation text
                if citation['type'] == 'file':
                    citation_text = f"📖 引用 {citation_num}\n\n"
                    citation_text += f"📄 文件：{citation['title']}\n\n"
                    citation_text += f"📝 內容：\n{citation['text']}"
                    if len(citation.get('text', '')) >= 500:
                        citation_text += "\n\n... (內容過長，已截斷)"
                elif citation['type'] == 'web':
                    citation_text = f"📖 引用 {citation_num}\n\n"
                    citation_text += f"🌐 來源：{citation['title']}\n"
                    citation_text += f"🔗 連結：{citation['uri']}"
                else:
                    citation_text = "無法顯示此引用。"

                reply_msg = TextSendMessage(text=citation_text)
                await line_bot_api.reply_message(event.reply_token, reply_msg)
                return
            else:
                reply_msg = TextSendMessage(text="找不到此引用，請重新查詢。")
                await line_bot_api.reply_message(event.reply_token, reply_msg)
                return
        except ValueError:
            pass  # Not a valid citation request, continue normal processing

    # Check if user wants to list files
    if is_list_files_intent(query):
        # Use File Manager Agent for conversational response
        agent = FileManagerAgent(store_name, store_name_cache)
        response_text = await agent.handle_list_files()
        reply_msg = TextSendMessage(text=response_text)
        await line_bot_api.reply_message(event.reply_token, reply_msg)
        return

    # Otherwise, query file search
    response_text, citations = await query_file_search(query, store_name)

    # Store citations in cache (limit to 3 for Quick Reply)
    if citations:
        citations_cache[store_name] = citations[:3]
        print(f"Stored {len(citations_cache[store_name])} citations for {store_name}")

    # Create Quick Reply buttons for citations
    quick_reply = None
    if citations:
        quick_reply_items = []
        for i, citation in enumerate(citations[:3], 1):  # Limit to 3 citations
            quick_reply_items.append(
                QuickReplyButton(action=MessageAction(
                    label=f"📖 引用{i}",
                    text=f"📖 引用{i}"
                ))
            )
        quick_reply = QuickReply(items=quick_reply_items)

    # Reply to user
    reply_msg = TextSendMessage(text=response_text, quick_reply=quick_reply)
    await line_bot_api.reply_message(event.reply_token, reply_msg)


@app.post("/")
async def handle_callback(request: Request):
    signature = request.headers["X-Line-Signature"]

    # Get request body as text
    body = await request.body()
    body = body.decode()

    try:
        events = parser.parse(body, signature)
    except InvalidSignatureError:
        raise HTTPException(status_code=400, detail="Invalid signature")

    for event in events:
        # Handle PostbackEvent (e.g., delete file button clicks)
        if isinstance(event, PostbackEvent):
            await handle_postback(event)
        # Handle MessageEvent
        elif isinstance(event, MessageEvent):
            if event.message.type == "text":
                # Process text message
                await handle_text_message(event, event.message)
            elif event.message.type == "file":
                # Process file message (upload to file search store)
                await handle_document_message(event, event.message)
            elif event.message.type == "image":
                # Process image message (analyze with Gemini vision)
                await handle_image_message(event, event.message)
            else:
                continue
        else:
            continue

    return "OK"


@app.on_event("shutdown")
async def shutdown_event():
    """Clean up resources on shutdown."""
    await client_session.close()