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

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

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 import 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 _load_environment(self):
        """加载环境变量"""
        load_dotenv()

    def _setup_knowledge_base(self) -> CombinedKnowledgeBase:
        """设置知识库"""
        local_pdf_kb = PDFKnowledgeBase(
            path=self.config.pdf_path,
            vector_db=LanceDb(
                table_name="database_fundamentals",
                uri=self.config.vector_db_path,
                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=self.config.vector_db_path,
                search_type=SearchType.vector,
                embedder=OpenAIEmbedder(id="text-embedding-3-small"),
            ),
        )
        knowledge_base.load()
        return knowledge_base

    def _setup_postgres_tools(self) -> MCPTools:
        """设置 PostgreSQL 工具"""
        server_params = StdioServerParameters(
            command="cmd",
            args=[
                "/c",
                "npx",
                "-y",
                "@modelcontextprotocol/server-postgres",
                self.config.db_connection,
            ],
            env={},
        )
        return MCPTools(server_params=server_params)

    def _create_agent(self) -> Agent:
        """创建智能体"""
        return Agent(
            model=DeepSeek(id=self.config.model_id),
            markdown=True,
            knowledge=self.knowledge_base,
            search_knowledge=True,
            show_tool_calls=True,
            tools=[self.postgres_tools],
        )

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

def main():
    """主入口函数"""
    optimizer = DatabaseOptimizer()
    optimizer.analyze_database("看下aq这个数据库")

if __name__ == "__main__":
    main()