"""数据库工具"""

import json
import os
import re
from typing import Any

from mcp.types import Tool, TextContent

from trailbase_mcp.client.client import TrailBaseClient


def get_db_tools() -> list[Tool]:
    """获取数据库相关工具"""
    return [
        Tool(
            name="db.query",
            description="执行 SQL 查询",
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {"type": "string", "description": "SQL 查询语句"},
                    "params": {
                        "type": ["array", "object"],
                        "description": "查询参数（可选）",
                    },
                    "timeout_ms": {"type": "integer", "description": "超时时间（毫秒）"},
                    "tx": {"type": "boolean", "description": "是否在事务中执行"},
                },
                "required": ["sql"],
            },
        ),
        Tool(
            name="db.explain",
            description="解释 SQL 查询计划",
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {"type": "string", "description": "SQL 查询语句"},
                },
                "required": ["sql"],
            },
        ),
        Tool(
            name="db.tables",
            description="列出数据库表",
            inputSchema={
                "type": "object",
                "properties": {
                    "like": {"type": "string", "description": "表名过滤模式（可选）"},
                },
            },
        ),
        Tool(
            name="db.table_info",
            description="获取表结构信息",
            inputSchema={
                "type": "object",
                "properties": {
                    "table": {"type": "string", "description": "表名"},
                },
                "required": ["table"],
            },
        ),
    ]


def _check_sql_safety(sql: str) -> tuple[bool, str | None]:
    """检查 SQL 语句的安全性
    
    Returns:
        (is_safe, error_message): 如果安全返回 (True, None)，否则返回 (False, 错误信息)
    """
    # 危险关键字列表（大写）
    dangerous_keywords = [
        "DROP", "ALTER", "ATTACH", "DETACH", "VACUUM",
        "PRAGMA", "CREATE", "DELETE", "INSERT", "UPDATE",
        "REPLACE", "TRUNCATE"
    ]
    
    # 检查是否允许危险操作（从环境变量读取）
    allow_dangerous = os.getenv("TRAILBASE_MCP_ALLOW_DANGEROUS_SQL", "false").lower() == "true"
    
    sql_upper = sql.upper().strip()
    
    # 允许 SELECT 和 EXPLAIN（只读操作）
    if sql_upper.startswith(("SELECT", "EXPLAIN", "WITH")):
        return (True, None)
    
    # 如果允许危险操作，跳过检查
    if allow_dangerous:
        return (True, None)
    
    # 检查是否包含危险关键字
    for keyword in dangerous_keywords:
        # 使用单词边界匹配，避免误判（如 "SELECT" 不会匹配 "SELECTED"）
        pattern = r'\b' + re.escape(keyword) + r'\b'
        if re.search(pattern, sql_upper):
            return (
                False,
                f"SQL contains dangerous keyword: {keyword}. "
                f"Set TRAILBASE_MCP_ALLOW_DANGEROUS_SQL=true to allow (not recommended for production)."
            )
    
    return (True, None)


async def handle_db_query(
    client: TrailBaseClient,
    sql: str,
    params: list[Any] | dict[str, Any] | None = None,
    timeout_ms: int | None = None,
    tx: bool | None = None,
) -> list[TextContent]:
    """处理 db.query 工具调用"""
    try:
        # SQL 安全检查
        is_safe, error_msg = _check_sql_safety(sql)
        if not is_safe:
            return [TextContent(
                type="text",
                text=json.dumps({"error": error_msg}, ensure_ascii=False)
            )]
        
        payload: dict[str, Any] = {"query": sql}
        if params is not None:
            payload["params"] = params
        if timeout_ms is not None:
            payload["timeout_ms"] = timeout_ms
        if tx is not None:
            payload["tx"] = tx

        result = client.http.post("/api/_admin/query", json=payload)
        return [
            TextContent(
                type="text",
                text=json.dumps(
                    {
                        "rows": result.get("rows", []),
                        "row_count": result.get("row_count", len(result.get("rows", []))),
                        "elapsed_ms": result.get("elapsed_ms", 0),
                    },
                    ensure_ascii=False
                ),
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}, ensure_ascii=False))]


async def handle_db_explain(
    client: TrailBaseClient,
    sql: str,
) -> list[TextContent]:
    """处理 db.explain 工具调用"""
    try:
        explain_sql = f"EXPLAIN QUERY PLAN {sql}"
        result = client.http.post("/api/_admin/query", json={"query": explain_sql})
        return [
            TextContent(
                type="text",
                text=json.dumps({"plan": result.get("rows", [])}, ensure_ascii=False),
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}, ensure_ascii=False))]


async def handle_db_tables(
    client: TrailBaseClient,
    like: str | None = None,
) -> list[TextContent]:
    """处理 db.tables 工具调用"""
    try:
        sql = "SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view')"
        params = None
        if like:
            sql += " AND name LIKE ?"
            params = [like]

        payload = {"query": sql}
        if params:
            payload["params"] = params
        result = client.http.post("/api/_admin/query", json=payload)
        rows = result.get("rows", [])
        columns = result.get("columns", [])

        tables = []
        for row in rows:
            # 解析行数据
            row_dict = {}
            if columns and len(row) == len(columns):
                for i, col in enumerate(columns):
                    col_name = col.get("name", f"col_{i}") if isinstance(col, dict) else (str(col) if col else f"col_{i}")
                    row_dict[col_name] = row[i] if i < len(row) else None
            else:
                row_dict = row if isinstance(row, dict) else {}
            
            table_name = row_dict.get("name")
            table_info: dict[str, Any] = {"name": table_name, "type": row_dict.get("type")}
            # 尝试获取行数
            try:
                if table_name:
                    count_result = client.http.post(
                        "/api/_admin/query",
                        json={"query": f'SELECT COUNT(*) as count FROM "{table_name}"'},
                    )
                    count_rows = count_result.get("rows", [])
                    if count_rows and len(count_rows) > 0:
                        table_info["row_count"] = count_rows[0][0] if count_rows[0] else 0
            except Exception:
                pass
            tables.append(table_info)

        return [TextContent(type="text", text=json.dumps({"tables": tables}, ensure_ascii=False))]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}, ensure_ascii=False))]


async def handle_db_table_info(
    client: TrailBaseClient,
    table: str,
) -> list[TextContent]:
    """处理 db.table_info 工具调用"""
    try:
        # 使用标准 SQL 查询获取表结构（从 sqlite_master）
        create_sql = f'SELECT sql FROM sqlite_master WHERE type="table" AND name="{table}"'
        create_result = client.http.post("/api/_admin/query", json={"query": create_sql})
        create_rows = create_result.get("rows", [])
        
        # 尝试使用 PRAGMA，如果失败则使用备用方法
        columns = []
        indexes = []
        
        try:
            # 获取列信息
            columns_sql = f'PRAGMA table_info("{table}")'
            columns_result = client.http.post("/api/_admin/query", json={"query": columns_sql})
            columns_rows = columns_result.get("rows", [])
            columns_cols = columns_result.get("columns", [])

            for row in columns_rows:
                # 解析行数据
                row_dict = {}
                if columns_cols and len(row) == len(columns_cols):
                    for i, col in enumerate(columns_cols):
                        col_name = col.get("name", f"col_{i}") if isinstance(col, dict) else (str(col) if col else f"col_{i}")
                        row_dict[col_name] = row[i] if i < len(row) else None
                else:
                    row_dict = row if isinstance(row, dict) else {}
                
                columns.append(
                    {
                        "name": row_dict.get("name"),
                        "type": row_dict.get("type"),
                        "not_null": bool(row_dict.get("notnull")),
                        "default": row_dict.get("dflt_value"),
                        "pk": bool(row_dict.get("pk")),
                    }
                )
        except Exception:
            # PRAGMA 失败，尝试从 CREATE TABLE 语句解析
            if create_rows and len(create_rows) > 0:
                create_sql_text = create_rows[0][0] if isinstance(create_rows[0], list) else create_rows[0]
                # 简单解析（这里可以改进）
                pass

        try:
            # 获取索引信息
            indexes_sql = f'PRAGMA index_list("{table}")'
            indexes_result = client.http.post("/api/_admin/query", json={"query": indexes_sql})
            indexes_rows = indexes_result.get("rows", [])
            indexes_cols = indexes_result.get("columns", [])

            for idx_row in indexes_rows:
                # 解析行数据
                idx_dict = {}
                if indexes_cols and len(idx_row) == len(indexes_cols):
                    for i, col in enumerate(indexes_cols):
                        col_name = col.get("name", f"col_{i}") if isinstance(col, dict) else (str(col) if col else f"col_{i}")
                        idx_dict[col_name] = idx_row[i] if i < len(idx_row) else None
                else:
                    idx_dict = idx_row if isinstance(idx_row, dict) else {}
                
                idx_name = idx_dict.get("name")
                idx_unique = bool(idx_dict.get("unique"))

                # 获取索引列
                idx_cols = []
                if idx_name:
                    idx_info_sql = f'PRAGMA index_info("{idx_name}")'
                    idx_info_result = client.http.post("/api/_admin/query", json={"query": idx_info_sql})
                    idx_info_rows = idx_info_result.get("rows", [])
                    idx_info_cols = idx_info_result.get("columns", [])

                    for r in idx_info_rows:
                        r_dict = {}
                        if idx_info_cols and len(r) == len(idx_info_cols):
                            for i, col in enumerate(idx_info_cols):
                                col_name = col.get("name", f"col_{i}") if isinstance(col, dict) else (str(col) if col else f"col_{i}")
                                r_dict[col_name] = r[i] if i < len(r) else None
                        else:
                            r_dict = r if isinstance(r, dict) else {}
                        idx_cols.append(r_dict.get("name"))
                
                indexes.append(
                    {
                        "name": idx_name,
                        "unique": idx_unique,
                        "columns": idx_cols,
                    }
                )
        except Exception:
            # PRAGMA 失败，索引列表为空
            pass

        return [
            TextContent(
                type="text",
                text=json.dumps({"columns": columns, "indexes": indexes}, ensure_ascii=False),
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({"error": str(e)}, ensure_ascii=False))]

