# ai_service.py (Instructs AI to execute tool calls without asking for confirmation)
import json
import os
import random
from datetime import datetime
import google.generativeai as genai
from google.api_core import exceptions as google_exceptions
from gradio_client import Client

# 從設定檔匯入金鑰和 URL
from config import MCP_SERVER_URL

# --- 1. API 金鑰健康檢查與輪替機制 ---
key1 = os.getenv("GEMINI_API_KEY")
key2 = os.getenv("GEMINI_API_KEY2")
all_keys = []
if key1 and "YOUR_GEMINI_API_KEY" not in key1: all_keys.append(key1)
if key2 and "YOUR_GEMINI_API_KEY" not in key2: all_keys.append(key2)

healthy_keys = []
if all_keys:
    print(f"--- Found {len(all_keys)} Gemini API Key(s). Starting validation... ---")
    for key in all_keys:
        try:
            genai.configure(api_key=key)
            test_model = genai.GenerativeModel('gemini-1.5-flash')
            test_model.generate_content("test", request_options={'timeout': 10})
            healthy_keys.append(key)
            print(f"--- Key ending in ...{key[-4:]} is VALID and added to the pool. ---")
        except Exception as e:
            print(f"--- Key ending in ...{key[-4:]} failed validation. Skipping. Error: {e} ---")
if not healthy_keys:
    print("--- CRITICAL: No valid Gemini API Keys found. AI Service will be disabled. ---")

# --- 2. 工具函式 (Tool Functions) ---
def call_mcp_earthquake_search(start_date: str, end_date: str, min_magnitude: float = 4.0, max_magnitude: float = 9.0) -> str:
    try:
        client = Client(src=MCP_SERVER_URL)
        result = client.predict(param_0=start_date, param_1="00:00:00", param_2=end_date, param_3="23:59:59", param_4=21.0, param_5=26.0, param_6=119.0, param_7=123.0, param_8=0.0, param_9=100.0, param_10=min_magnitude, param_11=max_magnitude, api_name="/gradio_fetch_and_plot_data")
        data = result[0].get('data', [])
        if not data: return "查詢完成，但未找到任何符合條件的地震資料。"
        headers = result[0].get('headers', [])
        return json.dumps([dict(zip(headers, row)) for row in data], indent=2, ensure_ascii=False)
    except Exception as e: return f"工具執行失敗，錯誤訊息: {e}"

def call_mcp_pws_search() -> str:
    try:
        client = Client("cwadayi/MCP-pws-running")
        result = client.predict(None, api_name="/get_disaster_warnings")
        return result[0] if isinstance(result, tuple) and len(result) > 0 else str(result)
    except Exception as e: return f"工具執行失敗，錯誤訊息: {e}"

# --- 3. 向 Gemini 定義工具 (Tool Declarations) ---
earthquake_search_tool_declaration = { "name": "call_earthquake_search_tool", "description": "從台灣中央氣象署的資料庫中搜尋地震事件。", "parameters": { "type": "OBJECT", "properties": { "start_date": { "type": "STRING", "description": "搜尋的開始日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間（例如：昨天、上個月、去年）或絕對時間（例如：2024年）來主動推斷此日期。" }, "end_date": { "type": "STRING", "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間或絕對時間來主動推斷此日期。" }, }, "required": ["start_date", "end_date"] } }
pws_search_tool_declaration = { "name": "call_mcp_pws_search", "description": "查詢最新的 PWS (Public Weather Service) 公共天氣服務發布情形。", "parameters": { "type": "OBJECT", "properties": {} } }
available_tools = { "call_earthquake_search_tool": call_mcp_earthquake_search, "call_mcp_pws_search": call_mcp_pws_search }

# --- 4. 主要的 AI 文字生成函式 ---
def generate_ai_text(user_prompt: str) -> str:
    if not healthy_keys:
        return "🤖 AI (Gemini) 服務錯誤：沒有任何有效的 API 金鑰可供使用。"
    
    try:
        selected_key = random.choice(healthy_keys)
        print(f"--- Handling request with key ending in: ...{selected_key[-4:]} ---")
        
        genai.configure(api_key=selected_key)
        
        # [*** 核心修正 ***]
        # 在系統指令中，明確要求 AI 在推斷出參數後「不要提問，直接執行」
        system_instruction = (
            "You are a helpful AI assistant. You must answer in Traditional Chinese."
            "When a user's query can be answered by a tool, you MUST infer the parameters from the query and call the tool immediately without asking for confirmation. Do not ask the user to confirm the parameters you have inferred."
            "When the user asks for the 'largest' or 'strongest' earthquake, you MUST assume they mean by magnitude. You should then analyze the JSON data returned by the tool, "
            "find the single entry with the highest 'magnitude' value, and present that specific earthquake's details as the answer."
        )
        
        model = genai.GenerativeModel(
            model_name="gemini-1.5-flash",
            tools=[earthquake_search_tool_declaration, pws_search_tool_declaration],
            system_instruction=system_instruction
        )
        
        chat = model.start_chat()
        response = chat.send_message(user_prompt)
        for part in response.parts:
            if part.function_call:
                function_call = part.function_call
                tool_function = available_tools.get(function_call.name)
                if not tool_function:
                    return f"錯誤：模型嘗試呼叫一個不存在的工具 '{function_call.name}'。"
                
                tool_result = tool_function(**dict(function_call.args))
                
                final_response = chat.send_message(
                    {"function_response": {"name": function_call.name, "response": {"result": tool_result}}}
                )
                return final_response.text
        
        return response.text

    except Exception as e:
        print(f"與 Gemini AI 互動時發生錯誤: {e}")
        return f"🤖 AI 服務發生錯誤: {e}"