#!/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
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


@dataclass
class DatabaseOptimizerConfig:
    """数据库优化器配置"""

    pdf_path: str = "D:\\Sources\\DONGJAK-TOOLS\\pdfs\\Database Fundamentals.pdf"
    db_connection: str = "postgresql://postgres:postgres@192.168.1.7:5432/postgres"
    model_id: str = "deepseek-chat"
    vector_db_path: str = "tmp/lancedb"


class DatabaseOptimizer:
    """PostgreSQL 数据库优化引擎"""

    def __init__(self, config: Optional[DatabaseOptimizerConfig] = None):
        self.config = config or DatabaseOptimizerConfig()
        self._load_environment()
        self.knowledge_base = self._setup_knowledge_base()
        self.postgres_tools = self._setup_postgres_tools()
        self.agent = self._create_agent()

    def _setup_postgres_tools(self) -> MCPTools:
        """设置 PostgreSQL 工具"""
        server_params = StdioServerParameters(
            command="cmd",
            args=[
                "/c",
                "npx",
                "-y",
                "@modelcontextprotocol/server-postgres",
                self.config.db_connection,
            ],
            env={},
        )
        # Create a client session to connect to the MCP server
        with stdio_client(server_params) as (read, write):
            with ClientSession(read, write) as session:
                agent = create_filesystem_agent(session)

                # Run the agent
                agent.print_response(message, stream=True)
        return MCPTools(server_params=server_params)

    def analyze_database(self, query: str, stream: bool = True):
        """分析数据库"""
        with self.postgres_tools:
            self.agent.print_response(query, stream=stream)


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"),
        ),
    )

    knowledge_base = CombinedKnowledgeBase(
        sources=[local_pdf_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()

    server_params = StdioServerParameters(
        command="cmd",
        args=[
            "/c",
            "npx",
            "-y",
            "@modelcontextprotocol/server-postgres",
            "postgresql://postgres:postgres@192.168.1.7:5432/postgres",
        ],
        env={},
    )
    # Create a client session to connect to the MCP server
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            postgres_tools = MCPTools(session=session)
            await postgres_tools.initialize()
            agent = Agent(
                model=DeepSeek(id="deepseek-chat"),
                markdown=True,
                knowledge=knowledge_base,
                search_knowledge=True,
                show_tool_calls=True,
                tools=[postgres_tools],
            )
            await agent.aprint_response("看下aq这个数据库", stream=True)


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