#!/usr/bin/env python3
"""
PostgreSQL 数据库优化工程师智能体

功能:
- 使用 DeepSeek 模型进行智能分析
- 集成 PostgreSQL MCP 服务器工具
- 加载数据库知识库
- 提供数据库优化建议
"""

import asyncio
import os
from typing import Optional
from dataclasses import dataclass
from dotenv import load_dotenv

from agno.agent import Agent, AgentMemory
from agno.models.deepseek import DeepSeek
from agno.knowledge.pdf import PDFKnowledgeBase
from agno.vectordb.lancedb import LanceDb, SearchType
from agno.embedder.openai import OpenAIEmbedder
from agno.knowledge.combined import CombinedKnowledgeBase
from agno.tools.mcp import MCPTools
from mcp.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters
from agno.knowledge.website import WebsiteKnowledgeBase
from agno.storage.agent.sqlite import SqliteAgentStorage
from agno.memory.db.sqlite import SqliteMemoryDb


async def main():
    load_dotenv()
    """设置知识库"""
    local_pdf_kb = PDFKnowledgeBase(
        path="D:\\Sources\\DONGJAK-TOOLS\\pdfs\\Database Fundamentals.pdf",
        vector_db=LanceDb(
            table_name="database_fundamentals",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )
    # Create Website knowledge base
    website_kb = WebsiteKnowledgeBase(
        urls=["https://www.lucidchart.com/blog/database-design-best-practices"],
        max_links=10,
        vector_db=LanceDb(
            table_name="website_documents",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )

    knowledge_base = CombinedKnowledgeBase(
        sources=[local_pdf_kb, website_kb],
        vector_db=LanceDb(
            table_name="combined_documents",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )
    knowledge_base.load()

    postgres_server_params = StdioServerParameters(
        command="cmd",
        args=[
            "/c",
            "npx",
            "-y",
            "@modelcontextprotocol/server-postgres",
            "postgresql://postgres:postgres@192.168.1.7:5432/aq",
        ],
        env={},
    )

    searxng_server_params = StdioServerParameters(
        command="cmd",
        args=[
            "/c",
            "npx",
            "-y",
            "https://github.com/ihor-sokoliuk/mcp-searxng.git",
        ],
        env={
            "SEARXNG_URL": "https://searchxng.ailoveworld.cn",
        },
    )
    # Create a client session to connect to the MCP server
    async with (
        MCPTools(server_params=postgres_server_params) as postgres_tools,
        MCPTools(server_params=searxng_server_params) as searxng_tools,
    ):
        agent = Agent(
            model=DeepSeek(id="deepseek-chat"),
            storage=SqliteAgentStorage(
                table_name="agent_sessions", db_file="tmp/agent_storage.db"
            ),
            memory=AgentMemory(
                db=SqliteMemoryDb(
                    table_name="agent_memory", db_file="tmp/agent_storage.db"
                ),
                create_user_memories=True,
                create_session_summary=True,
            ),
            # Set add_history_to_messages=true to add the previous chat history to the messages sent to the Model.
            add_history_to_messages=True,
            # Number of historical responses to add to the messages.
            num_history_responses=3,
            markdown=True,
            knowledge=knowledge_base,
            search_knowledge=True,
            show_tool_calls=True,
            tools=[postgres_tools, searxng_tools],
            session_id="1",
            # Enable the agent to read the chat history
            read_chat_history=True,
        )
        # await agent.aprint_response("我现在想要记录用户对智能体和课程的使用权限及使用情况，需要如何设计表结构", stream=True)
        await agent.aprint_response("调用get_chat_history", stream=True)
        # await agent.aprint_response("帮我分析一下aq.public数据库,并给出优化建议", stream=True)
        # await agent.aprint_response("阅读下 https://www.lucidchart.com/blog/database-design-best-practices 这篇文章", stream=True)


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