import asyncio
import logging
import os
from typing import Optional, Dict, Any, List
import psycopg
from psycopg.rows import dict_row
from mcp.server.fastmcp import FastMCP

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化MCP server
mcp = FastMCP(name="PostgreSQL Database Server")

# 数据库连接配置
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://user:password@localhost:5432/dbname")

def get_db_connection():
    """获取数据库连接"""
    try:
        return psycopg.connect(DATABASE_URL, row_factory=dict_row)
    except Exception as e:
        logger.error(f"Database connection failed: {e}")
        raise

@mcp.tool()
def list_tables(schema: Optional[str] = None) -> Dict[str, Any]:
    """List user tables (exclude system schemas)."""
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                if schema:
                    sql = """
                        SELECT table_name 
                        FROM information_schema.tables 
                        WHERE table_schema = %s 
                        AND table_type = 'BASE TABLE'
                        ORDER BY table_name
                    """
                    cur.execute(sql, (schema,))
                else:
                    sql = """
                        SELECT table_name 
                        FROM information_schema.tables 
                        WHERE table_schema NOT IN ('information_schema', 'pg_catalog', 'pg_toast')
                        AND table_type = 'BASE TABLE'
                        ORDER BY table_schema, table_name
                    """
                    cur.execute(sql)
                
                tables = [row['table_name'] for row in cur.fetchall()]
                return {"tables": tables}
    except Exception as e:
        logger.error(f"Error listing tables: {e}")
        return {"error": str(e), "tables": []}

@mcp.tool()
def execute_sql(sql: str) -> Dict[str, Any]:
    """Execute arbitrary SQL. Returns rows for SELECT else rowCount."""
    if not sql or not isinstance(sql, str):
        return {"error": "Invalid SQL statement"}
    
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(sql)
                
                # 判断是否为查询语句
                if sql.strip().upper().startswith(('SELECT', 'WITH', 'SHOW', 'DESCRIBE', 'EXPLAIN')):
                    rows = cur.fetchall()
                    return {"rows": rows, "rowCount": len(rows)}
                else:
                    # 对于INSERT/UPDATE/DELETE等语句
                    conn.commit()
                    return {"rowCount": cur.rowcount, "message": "Query executed successfully"}
                    
    except Exception as e:
        logger.error(f"Error executing SQL: {e}")
        return {"error": str(e)}


async def main():
    """Start the MCP server with streamable transport."""
    logger.info("Starting PostgreSQL MCP Server")
    
    # Start the streamable HTTP server
    await mcp.run_streamable_http_async()

def sync_main():
    """Synchronous entry point for the server."""
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Server shutdown requested")
    except Exception as e:
        logger.error(f"Server error: {e}")
        raise

if __name__ == "__main__":
    sync_main()