"""
MCP MySQL Server
主逻辑文件
"""

from typing import List, Dict, Any
import pymysql
from mcp.server.fastmcp import FastMCP
import os

# 读取数据库配置（可用.env或环境变量注入）
DB_USER = os.environ.get("DB_USER", "readonly_user")
DB_PASS = os.environ.get("DB_PASS", "readonly_pass")
DB_HOST = os.environ.get("DB_HOST", "localhost")
DB_PORT = int(os.environ.get("DB_PORT", "3306"))
DB_NAME = os.environ.get("DB_NAME", "test_db")

# MCP server 初始化
mcp = FastMCP("mysql-stdio")

# MCP资源：读取知识文件
@mcp.resource("file://knowledge.txt")
def read_knowledge() -> str:
	try:
		with open("knowledge.txt", "r", encoding="utf-8") as f:
			return f.read()
	except Exception as e:
		return f"[ERROR] Cannot read knowledge.txt: {e}"
    
# 数据库操作函数
def query_mysql(query: str, params=None) -> List[Dict[str, Any]]:
    """执行MySQL查询，让异常自然传播"""
    conn = pymysql.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASS,
        database=DB_NAME,
        charset="utf8mb4",
        autocommit=True,
        cursorclass=pymysql.cursors.DictCursor
    )
    
    try:
        with conn.cursor() as cursor:
            cursor.execute(query, params)
            result = cursor.fetchall()
        return result
    finally:
        conn.close()

# MCP资源：列出所有表
@mcp.resource()
async def list_tables() -> List[dict]:
    """列出数据库中所有表"""
    result = query_mysql("SHOW TABLES")
    if not result:
        return []
    
    key = next(iter(result[0]))
    return [
        {
            "id": row[key],
            "name": row[key],
            "description": f"MySQL table: {row[key]}",
            "mimeType": "application/sql-table",
        }
        for row in result
    ]

# MCP工具：表结构
@mcp.tool()
async def describe_table(table_name: str) -> str:
    """获取表结构信息"""
    if not table_name or not table_name.strip():
        raise ValueError("表名不能为空")
    
    result = query_mysql(f"DESCRIBE `{table_name}`")
    if not result:
        return f"表 '{table_name}' 不存在或查询失败"
    
    lines = [f"表 `{table_name}` 的结构:"]
    for col in result:
        lines.append(
            f"  {col['Field']} | {col['Type']} | Null: {col['Null']} | Key: {col['Key']} | Default: {col['Default']}"
        )
    return "\n".join(lines)

# MCP工具：通用数据查询
@mcp.tool()
async def query_table(
    table_name: str,
    limit: int = 10,
    offset: int = 0,
    where: str = "",
    fields: str = "",
    order_by: str = ""
) -> str:
    """查询表数据"""
    if not table_name or not table_name.strip():
        raise ValueError("表名不能为空")
    
    if limit <= 0 or limit > 1000:
        raise ValueError("limit必须在1-1000之间")
    
    if offset < 0:
        raise ValueError("offset不能为负数")
    
    select_fields = fields if fields else "*"
    sql = f"SELECT {select_fields} FROM `{table_name}`"
    if where:
        sql += f" WHERE {where}"
    if order_by:
        sql += f" ORDER BY {order_by}"
    sql += f" LIMIT %s OFFSET %s"
    params = (limit, offset)
    
    result = query_mysql(sql, params)
    if not result:
        return f"表 '{table_name}' 中没有找到数据或查询失败"
    
    headers = list(result[0].keys())
    rows = [list(row.values()) for row in result]
    md = ["| " + " | ".join(headers) + " |", "|" + "|".join(["---"] * len(headers)) + "|"]
    for row in rows:
        md.append("| " + " | ".join([str(v) if v is not None else "" for v in row]) + " |")
    return f"查询 `{table_name}` 结果 (limit={limit}, offset={offset}):\n" + "\n".join(md)
