#!/usr/bin/env python3
"""
MySQL 协议数据库 MCP Server - 第一版

- 连接与元数据：list_databases, list_tables, describe_table
- SQL 分析：analyze_sql（支持多条，返回列表）
- SQL 执行：execute_sql（仅单条，DDL 不执行，写操作需 confirm）
- 多环境：connection 参数切换 profile
"""

from __future__ import annotations

import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Annotated, Optional

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field

from .client import (
    connect,
    describe_table as _describe_table,
    execute_select,
    get_connection_config,
    list_databases as _list_databases,
    list_tables as _list_tables,
)
from .config import (
    ConnectionConfig,
    SafetyConfig,
    load_connection_map,
    load_safety_config,
)
from .sql_analysis import (
    analyze_many,
    analyze_one,
    ensure_single_statement,
    StatementAnalysis,
)


# ============== 配置加载 ==============

DEFAULT_CONN, CONNECTION_MAP, DEFAULT_PROFILE_NAME = load_connection_map()
SAFETY = load_safety_config()


@dataclass
class DBState:
    """MCP 状态"""
    default_connection: ConnectionConfig
    connection_map: dict[str, ConnectionConfig]
    default_profile_name: str
    safety: SafetyConfig


@asynccontextmanager
async def lifespan(app: FastMCP) -> AsyncIterator[DBState]:
    """服务器生命周期"""
    state = DBState(
        default_connection=DEFAULT_CONN,
        connection_map=CONNECTION_MAP,
        default_profile_name=DEFAULT_PROFILE_NAME,
        safety=SAFETY,
    )
    print("Database MCP Server 已启动", file=sys.stderr)
    print(f"  默认 profile: {DEFAULT_PROFILE_NAME}", file=sys.stderr)
    print(f"  可用 profile: {', '.join(sorted(CONNECTION_MAP))}", file=sys.stderr)
    yield state
    print("Database MCP Server 已关闭", file=sys.stderr)


mcp = FastMCP("MySQL数据库", lifespan=lifespan)


# ============== 数据模型 ==============

class ConnectionItem(BaseModel):
    profile: str = Field(description="profile 名，即 connection 参数值")
    description: Optional[str] = Field(default=None, description="可选说明/别名，便于根据用户自然语言匹配")
    databases: Optional[list[str]] = Field(default=None, description="该连接聚焦的数据库列表，不设则无限制")


class ConnectionList(BaseModel):
    connections: list[ConnectionItem] = Field(description="可用连接列表")
    default: str = Field(description="默认 profile 名")


class DatabaseList(BaseModel):
    project: str = Field(description="profile 名")
    databases: list[str] = Field(description="数据库名列表")
    count: int = Field(description="数量")


class TableList(BaseModel):
    database: str = Field(description="数据库名")
    tables: list[dict] = Field(description="表列表，含 name/rows_approx/comment")
    count: int = Field(description="数量")


class TableDescribe(BaseModel):
    database: str = Field(description="数据库名")
    table: str = Field(description="表名")
    columns: list[dict] = Field(description="列信息")
    rows_approx: Optional[int] = Field(default=None, description="近似行数")
    size_kb: Optional[int] = Field(default=None, description="表大小(KB)")


class AnalysisItem(BaseModel):
    sql: str = Field(description="SQL 片段")
    stype: str = Field(description="类型: DDL/SELECT/INSERT/UPDATE/DELETE/OTHER")
    first_keyword: str = Field(description="首关键字")
    need_confirm: bool = Field(description="是否需要 confirm 后执行")
    message: str = Field(description="说明")


class AnalyzeResult(BaseModel):
    statements: list[AnalysisItem] = Field(description="每条的分析结果")
    count: int = Field(description="语句条数")


class ExecuteResult(BaseModel):
    stype: str = Field(description="语句类型")
    rows: Optional[list[dict]] = Field(default=None, description="SELECT 结果行")
    rows_affected: Optional[int] = Field(default=None, description="写操作影响行数")
    truncated: bool = Field(default=False, description="是否被截断")
    message: str = Field(default="", description="附加说明")


def _get_state() -> DBState:
    return mcp.get_context().request_context.lifespan_context


def _conn(profile: Optional[str]) -> ConnectionConfig:
    state = _get_state()
    return get_connection_config(profile, state.default_connection, state.connection_map)


def _check_database_allowed(config: ConnectionConfig, database: str) -> None:
    """未配置 databases 则允许所有库；配置为空数组则不允许任何库；配置为非空数组则仅允许列表中的库。"""
    dbs = getattr(config, "databases", None)
    if dbs is None:
        return  # 未配置：默认所有库
    if len(dbs) == 0:
        raise ValueError("该连接已配置为不允许访问任何数据库（databases 为空数组）。")
    if database not in dbs:
        raise ValueError(
            f"该连接仅允许访问以下数据库: {', '.join(dbs)}；当前指定了 '{database}'。"
        )


# ============== Tools ==============

@mcp.tool()
def list_connections() -> ConnectionList:
    """列出所有可用连接（profile 名及可选说明）。便于根据用户自然语言（如「测试1」「生产」）选择 connection 参数。"""
    state = _get_state()
    items = [
        ConnectionItem(
            profile=k,
            description=getattr(state.connection_map[k], "description", None),
            databases=getattr(state.connection_map[k], "databases", None),
        )
        for k in sorted(state.connection_map)
    ]
    return ConnectionList(connections=items, default=state.default_profile_name)


@mcp.tool()
def list_databases(
    connection: Annotated[
        Optional[str],
        Field(description="profile 名，如 dev/prod；不填使用默认连接"),
    ] = None,
) -> DatabaseList:
    """列出当前连接下所有数据库。若该连接配置了聚焦数据库（databases 数组），则只返回这些库。"""
    state = _get_state()
    config = _conn(connection)
    with connect(config, state.safety) as conn:
        dbs = _list_databases(conn, getattr(config, "databases", None))
    profile = (connection or "").strip() or state.default_profile_name
    return DatabaseList(project=profile, databases=dbs, count=len(dbs))


@mcp.tool()
def list_tables(
    database: Annotated[str, Field(description="数据库名")],
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> TableList:
    """列出指定库下的表（含可选行数/备注）。若该连接配置了聚焦数据库，仅能查这些库。"""
    state = _get_state()
    config = _conn(connection)
    _check_database_allowed(config, database)
    with connect(config, state.safety) as conn:
        tables = _list_tables(conn, database)
    return TableList(
        database=database,
        tables=tables,
        count=len(tables),
    )


@mcp.tool()
def describe_table(
    database: Annotated[str, Field(description="数据库名")],
    table: Annotated[str, Field(description="表名")],
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> TableDescribe:
    """查看单表结构：列名、类型、主键、可选行数/大小。若该连接配置了聚焦数据库，仅能查这些库。"""
    state = _get_state()
    config = _conn(connection)
    _check_database_allowed(config, database)
    with connect(config, state.safety) as conn:
        info = _describe_table(conn, database, table)
    return TableDescribe(
        database=info["database"],
        table=info["table"],
        columns=info["columns"],
        rows_approx=info.get("rows_approx"),
        size_kb=info.get("size_kb"),
    )


@mcp.tool()
def analyze_sql(
    sql: Annotated[str, Field(description="SQL 文本，可含多条（; 分隔）")],
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> AnalyzeResult:
    """解析并分析 SQL，不执行。支持多条，返回每条的类型与是否需确认；DDL 会提示手动执行。"""
    _conn(connection)  # 校验 profile
    state = _get_state()
    analyses = analyze_many(sql, state.safety.read_only)
    items = [
        AnalysisItem(
            sql=a.sql,
            stype=a.stype,
            first_keyword=a.first_keyword,
            need_confirm=a.need_confirm,
            message=a.message,
        )
        for a in analyses
    ]
    return AnalyzeResult(statements=items, count=len(items))


@mcp.tool()
def execute_sql(
    sql: Annotated[str, Field(description="单条 SQL")],
    confirm: Annotated[
        Optional[bool],
        Field(description="写操作必须为 true 才执行"),
    ] = None,
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
    limit: Annotated[Optional[int], Field(description="SELECT 限制行数")] = None,
    offset: Annotated[Optional[int], Field(description="SELECT 偏移")] = None,
) -> ExecuteResult:
    """执行单条 SQL。DDL 不执行仅提示；SELECT 应用行数限制；写操作需 confirm=true。"""
    state = _get_state()
    single = ensure_single_statement(sql)
    analysis = analyze_one(single, state.safety.read_only)

    if analysis.stype == "DDL":
        return ExecuteResult(
            stype="DDL",
            message="DDL 请在数据库客户端中手动执行，本 MCP 不执行 DDL。",
        )

    if analysis.stype == "SELECT":
        config = _conn(connection)
        with connect(config, state.safety) as conn:
            rows, truncated = execute_select(
                conn,
                single,
                state.safety.max_select_rows,
                limit=limit,
                offset=offset,
            )
        return ExecuteResult(
            stype="SELECT",
            rows=rows,
            truncated=truncated,
            message="结果已截断，请加 LIMIT 或分页查询。" if truncated else "",
        )

    if analysis.stype in ("INSERT", "UPDATE", "DELETE"):
        if state.safety.read_only:
            return ExecuteResult(
                stype=analysis.stype,
                message="当前为只读模式，拒绝执行写操作。",
            )
        if not confirm:
            return ExecuteResult(
                stype=analysis.stype,
                message="写操作需传入 confirm=true 后执行。",
            )
        config = _conn(connection)
        with connect(config, state.safety) as conn:
            with conn.cursor() as cur:
                cur.execute(single)
                affected = cur.rowcount
            conn.commit()
        return ExecuteResult(
            stype=analysis.stype,
            rows_affected=affected,
            message=f"已执行，影响 {affected} 行。",
        )

    return ExecuteResult(
        stype=analysis.stype,
        message=analysis.message or "未识别的语句类型，未执行。",
    )


# ============== Resources ==============

def _get_state_safe() -> DBState | None:
    """获取 state，资源在无请求上下文时可能为 None"""
    try:
        return mcp.get_context().request_context.lifespan_context
    except Exception:
        return None


@mcp.resource("db://config")
def resource_config() -> str:
    """当前连接与安全策略摘要（不暴露密码）。含各 profile 及可选说明，便于智能体根据用户自然语言选连接。"""
    state = _get_state_safe()
    if not state:
        return "# 配置信息需在会话中获取。"
    s = state.safety
    lines = [
        "# 数据库 MCP 配置",
        "",
        "## 默认 profile",
        f"- {state.default_profile_name}",
        "",
        "## 可用 profile（调用工具时传 connection=\"<profile>\"）",
    ]
    for k in sorted(state.connection_map):
        cfg = state.connection_map[k]
        desc = getattr(cfg, "description", None) or ""
        dbs = getattr(cfg, "databases", None)
        parts = [f"- **{k}**"]
        if desc:
            parts.append(desc)
        if dbs:
            parts.append(f"聚焦库: {', '.join(dbs)}")
        lines.append(" ".join(parts) if len(parts) > 1 else parts[0])
    lines.extend([
        "",
        "## 安全与行为",
        f"- 只读模式: {s.read_only}",
        f"- 单次 SELECT 最大行数: {s.max_select_rows}",
        f"- 查询超时(秒): {s.query_timeout_seconds}",
        f"- 慢查询阈值(秒): {s.slow_query_seconds}",
    ])
    return "\n".join(lines)


@mcp.resource("db://databases")
def resource_databases() -> str:
    """数据库列表（默认连接）"""
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"
    config = state.default_connection
    with connect(config, state.safety) as conn:
        dbs = _list_databases(conn)
    return "\n".join([f"- {d}" for d in dbs]) if dbs else "(无)"


@mcp.resource("db://tables/{database}")
def resource_tables(database: str) -> str:
    """指定库下的表列表"""
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"
    config = state.default_connection
    with connect(config, state.safety) as conn:
        tables = _list_tables(conn, database)
    lines = [f"- {t.get('name', t)} (行数约: {t.get('rows_approx', '?')})" for t in tables]
    return "\n".join(lines) if lines else "(无)"


@mcp.resource("db://table/{database}/{table}")
def resource_table(database: str, table: str) -> str:
    """单表结构摘要"""
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"
    config = state.default_connection
    with connect(config, state.safety) as conn:
        info = _describe_table(conn, database, table)
    lines = [
        f"# {database}.{table}",
        f"- 近似行数: {info.get('rows_approx', '?')}",
        f"- 大小(KB): {info.get('size_kb', '?')}",
        "",
        "## 列",
    ]
    for c in info.get("columns", []):
        lines.append(f"- {c.get('COLUMN_NAME')} {c.get('DATA_TYPE')} {c.get('COLUMN_KEY') or ''}")
    return "\n".join(lines)


# ============== Prompts ==============

@mcp.prompt()
def ask_user_which_connection_prompt() -> str:
    """连接选择规则：一开始无配置需用户选择，选后记住且不再提示；仅切换时再列出让用户选"""
    return """数据库连接的选用规则（请在对话中记住「当前使用的 connection」）：

1. **一开始没有任何已选配置**：在使用任何数据库工具前，必须先调用 list_connections，将可用连接列表展示给用户，并询问「请选择要使用的连接（如 test_0 / uat_0 / prod_0，或说“用测试”“用生产”）」。等用户选择或说出环境后，将该 connection 记作「当前使用的连接」，并在本次及后续所有数据库工具调用中默认使用该 connection，**不再重复提示**。

2. **仅当用户明确要切换环境时**（如说「切换」「换一个环境」「用测试」「连 UAT」）：再次调用 list_connections，将列表展示给用户并让其选择；用户选择后更新「当前使用的连接」，后续默认使用新 connection，直至用户再次要求切换。

3. 总结：无配置 → 列出并让用户选一次；选后记住，后续不提示；只有用户要切换时才再次列出并让用户选。
"""


@mcp.prompt()
def query_safely_prompt() -> str:
    """安全查询与执行习惯"""
    return """使用本数据库 MCP 时请遵守：

1. 优先用 SELECT 探索数据，对大表或未知行数查询务必加 LIMIT。
2. 写操作（INSERT/UPDATE/DELETE）前先使用 analyze_sql 查看，执行时传入 confirm=true。
3. DDL（CREATE/ALTER/DROP 等）不在本 MCP 执行，请在数据库客户端中手动执行。
4. 每次只执行一条 SQL；多步操作请多次调用 execute_sql。
"""


@mcp.prompt()
def explore_schema_prompt(
    database: Annotated[str, Field(description="数据库名")] = "",
    connection: Annotated[str, Field(description="profile 名")] = "prod",
) -> str:
    """探索 Schema 再写 SQL"""
    return f"""在编写或分析 SQL 前，请先获取相关表结构：

1. 使用 list_tables(database="{database or '<数据库名>'}", connection="{connection}") 查看表列表。
2. 使用 describe_table(database="...", table="...", connection="{connection}") 或读取 db://table/{{database}}/{{table}} 查看列与类型。
3. 再根据结构生成或分析 SQL。
"""


@mcp.prompt()
def switch_connection_prompt(
    profile: Annotated[str, Field(description="profile 名，如 dev/prod/test1")] = "prod",
) -> str:
    """自然语言切换环境时，后续工具调用应带上的 connection 参数"""
    return f"""用户希望使用 **{profile}** 环境时，后续调用数据库相关工具请传入 connection="{profile}"。
例如：list_databases(connection="{profile}")、execute_sql(sql="...", connection="{profile}")。
"""


@mcp.prompt()
def use_connection_from_natural_language_prompt() -> str:
    """根据用户选择或自然语言确定 connection；切换时也需列出让用户选"""
    return """确定数据库连接（connection）时：

1. **首次使用**：当前还没有「已选连接」时，先 list_connections，把列表展示给用户并让其选择；根据用户选择或说法（如「用测试」「prod」）映射到 profile 名，记作当前连接，后续一律用该 connection，不再问。
2. **用户要求切换**：当用户说「切换环境」「换一个」「用 UAT」等时，再次 list_connections，展示给用户并让其选择，选后更新当前连接。
3. 将用户说法映射到 profile：如「用测试」「测试环境」→ test_0/test_finance 等，「生产」「prod」→ prod_0 等；不确定时以 list_connections 返回的 profile 名为准。
"""


# ============== 入口 ==============

def main() -> None:
    mcp.run()


if __name__ == "__main__":
    main()
