import gradio as gr
import os
import time

# ------------------------------------------------------------------
# 1. 加载环境变量 (在Hugging Face Spaces中从Secrets加载)
# ------------------------------------------------------------------
# 使用os.environ.get()来安全地获取，如果没有设置，可以给一个默认值或报错
# 这里假设你在Hugging Face Spaces的Secrets中设置了'GEMINI_API_KEY'
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
    print("警告：未找到 GEMINI_API_KEY。请在Hugging Face Spaces的Secrets中设置它。")
    # 为了本地测试，可以从.env文件加载(需要安装python-dotenv)
    # from dotenv import load_dotenv
    # load_dotenv()
    # api_key = os.environ.get("GEMINI_API_KEY")

# ------------------------------------------------------------------
# 2. 初始化后端 (这是整个系统的启动点)
# ------------------------------------------------------------------
print("--- 正在启动 AI 助理系统 ---")

# 导入我们的核心模块
# 使用try...except来处理可能的导入错误，这在调试时很有用
try:
    from database.setup import initialize_system
    from core.agent import SmartAIAgent

    print("核心模块导入成功。")
except ImportError as e:
    print(f"导入模块时出错: {e}")
    print("请确保所有项目文件都已正确放置。")
    # 如果导入失败，系统无法运行，这里可以抛出异常或退出
    raise

# 执行一次性的系统初始化（创建数据库、向量索引、加载工具等）
# 这个函数应该被设计为幂等的，即多次运行不会产生副作用
try:
    # initialize_system() 函数将负责所有数据库和向量库的设置
    # 它会返回一个已注册工具的列表或其他必要信息
    registered_tools, tool_recommender = initialize_system()
    print("系统数据库和工具推荐器初始化完成。")

    # 创建 AI 智能体实例，并将推荐器和工具传递给它
    # Agent需要知道所有可以执行的工具函数
    agent = SmartAIAgent(
        tool_recommender=tool_recommender,
        registered_tools=registered_tools,
        api_key=api_key,
    )
    print("AI 智能体核心已成功创建。")
except Exception as e:
    print(f"系统初始化过程中发生严重错误: {e}")
    agent = None  # 标记Agent不可用
    # 在Gradio界面上可以显示错误信息

print("--- AI 助理系统已准备就绪 ---")

# ------------------------------------------------------------------
# 3. Gradio 事件处理函数
# ------------------------------------------------------------------


def handle_user_message(user_input, history):
    """
    当用户发送消息时，此函数首先被调用。
    它将用户的消息添加到聊天历史记录中。
    """
    if not user_input.strip():
        # 如果用户输入为空，不做任何事
        return "", history
    # 将用户消息和一条空的机器人消息占位符添加到历史
    history.append((user_input, None))
    # 返回空字符串以清空输入框，并返回更新后的历史记录
    return "", history


def generate_bot_response(history):
    """
    此函数以流式方式生成机器人的响应。
    它调用Agent核心来处理最新的用户消息。
    """
    if agent is None:
        # 如果Agent初始化失败，返回错误信息
        history[-1][1] = "抱歉，AI助理系统初始化失败，无法提供服务。"
        yield history
        return

    # 获取最新的用户问题
    user_question = history[-1][0]

    # 初始化一个空的机器人消息
    bot_message = ""
    history[-1][1] = bot_message

    try:
        # 调用Agent的流式处理方法
        # agent.stream_run() 应该是一个生成器，逐步yield出思考过程和最终答案
        for chunk in agent.stream_run(user_question):
            # 将每个新的文本块追加到机器人消息中
            bot_message += chunk
            # 更新历史记录中最后一个元组的机器人部分
            history[-1][1] = bot_message
            # yield更新后的历史记录，Gradio会用它来刷新界面
            yield history
            # (可选) 增加一个微小的延迟，让流式效果更明显
            time.sleep(0.01)

    except Exception as e:
        # 如果在处理过程中发生错误，将错误信息显示给用户
        error_message = f"\n\n抱歉，处理您的请求时发生了错误：\n`{str(e)}`"
        history[-1][1] += error_message
        yield history


# ------------------------------------------------------------------
# 4. 创建 Gradio 界面
# ------------------------------------------------------------------

# 自定义CSS来美化界面
custom_css = """
/* 简单的CSS自定义，让界面更好看 */
#chatbot .user {
    background-color: #E0F7FA; /* 浅青色背景 */
}
#chatbot .bot {
    background-color: #F1F8E9; /* 浅绿色背景 */
}
"""

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="lime"),
    css=custom_css,
    title="智能 AI 助理",
) as demo:
    gr.Markdown(
        """
        # 🚀 智能 AI 助理 Demo
        ### (LangChain + LlamaIndex + Gemini)
        这是一个演示如何结合 LangChain (智能体编排) 和 LlamaIndex (工具检索) 构建高级AI助理的Demo。
        - **提问:** 在下面的文本框中输入你的问题。
        - **观察:** 观察AI的思考过程，包括它如何推荐、选择和调用工具。
        """
    )

    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        label="聊天窗口",
        bubble_full_width=False,
        height=650,
        avatar_images=(
            None,
            "https://raw.githubusercontent.com/gradio-app/gradio/main/guides/assets/logo.png",
        ),  # 机器人用Gradio logo
    )

    with gr.Row():
        text_input = gr.Textbox(
            scale=4,
            show_label=False,
            placeholder="例如: '苹果公司(AAPL)今天的股价是多少？' 或 '关于AI的最新进展有什么新闻？'",
            container=False,
        )
        submit_button = gr.Button("发送", variant="primary", scale=1, min_width=150)

    # 示例问题，方便用户快速体验
    gr.Examples(
        examples=[
            "苹果公司(AAPL)的股价是多少？",
            "关于AI驱动的药物发现有什么最新新闻？",
            "你好，你能做什么？",
            "用Python写一个快速排序算法",  # 测试不使用工具的场景
        ],
        inputs=text_input,
    )

    # 定义事件的触发流程
    # 当用户提交输入时（点击按钮或按回车）
    submit_event = text_input.submit(
        fn=handle_user_message,
        inputs=[text_input, chatbot],
        outputs=[text_input, chatbot],
        queue=False,  # 立即执行，不排队
    ).then(
        fn=generate_bot_response,
        inputs=[chatbot],
        outputs=[chatbot],
    )

    submit_button.click(
        fn=handle_user_message,
        inputs=[text_input, chatbot],
        outputs=[text_input, chatbot],
        queue=False,
    ).then(
        fn=generate_bot_response,
        inputs=[chatbot],
        outputs=[chatbot],
    )

# ------------------------------------------------------------------
# 5. 启动应用
# ------------------------------------------------------------------
if __name__ == "__main__":
    # 使用 .queue() 来允许多个用户同时使用，这在Hugging Face Spaces上是推荐做法
    demo.queue()
    # .launch() 会启动Web服务器
    # 在Hugging Face Spaces上，它会自动找到并运行这个
    demo.launch(debug=True)  # debug=True可以在本地看到更详细的日志