import os
import asyncio

from base.agent import HuiYuBaseAgent
from google.adk.models.lite_llm import LiteLlm
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types

# ============================================================
# MiniMax 模型配置（从 .env 文件读取）
# ============================================================
MINIMAX_API_KEY = os.getenv("MINIMAX_API_KEY")
MINIMAX_API_BASE = os.getenv("MINIMAX_API_BASE")
MINIMAX_MODEL = os.getenv("MINIMAX_MODEL")

# ============================================================
# 创建 LiteLlm 模型适配器
# ============================================================
model = LiteLlm(
    model=MINIMAX_MODEL,
    api_base=MINIMAX_API_BASE,
    api_key=MINIMAX_API_KEY,
)

# ============================================================
# 定义 Agent（继承 HuiYuBaseAgent，自动获得防注入 + 日志能力）
# ============================================================
root_agent = HuiYuBaseAgent(
    name="minimax_agent",
    model=model,
    description="一个使用 MiniMax 模型的智能助手，能够回答用户的各种问题。",
    instruction="你是一个乐于助人的中文 AI 助手，由 MiniMax 模型驱动。请用中文回答用户的问题，回答要准确、简洁、友好。",
)

# ============================================================
# 运行入口（用于 adk run / adk web）
# ============================================================
APP_NAME = "慧遇agent app"
USER_ID = "user_001"
SESSION_ID = "session_001"


async def main():
    """主函数：创建会话并运行 Agent 对话"""
    # 1. 创建会话服务
    session_service = InMemorySessionService()

    # 2. 创建会话
    session = await session_service.create_session(
        app_name=APP_NAME,
        user_id=USER_ID,
        session_id=SESSION_ID,
    )
    print(f"会话已创建: App='{APP_NAME}', User='{USER_ID}', Session='{SESSION_ID}'")

    # 3. 创建 Runner
    runner = Runner(
        agent=root_agent,
        app_name=APP_NAME,
        session_service=session_service,
    )
    print(f"Runner 已创建，Agent: '{runner.agent.name}'")

    # 4. 交互式对话循环
    print("\n===== MiniMax Agent 对话 (输入 'exit' 退出) =====")
    while True:
        query = input("\n[user]: ").strip()
        if query.lower() in ("exit", "quit", "退出"):
            print("再见！")
            break
        if not query:
            continue

        # 构造 ADK 消息
        content = types.Content(role="user", parts=[types.Part(text=query)])

        # 调用 Agent
        final_response_text = ""
        async for event in runner.run_async(
            user_id=USER_ID,
            session_id=SESSION_ID,
            new_message=content,
        ):
            if event.is_final_response():
                if event.content and event.content.parts:
                    final_response_text = event.content.parts[0].text

        print(f"\n[agent]: {final_response_text}")


if __name__ == "__main__":
    asyncio.run(main())