#!/usr/bin/env python3
"""
Redis MCP Server - 单机 + 集群

- 元数据：list_connections, list_keys, info
- 命令分析：analyze_command（不执行）
- 执行：execute_command（单条，由大模型产生；MCP 只做安全检查与执行）
"""

from __future__ import annotations

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

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

from .client import (
    connect,
    execute_command as _execute_command,
    get_connection_config,
    info_sections,
    list_keys as _list_keys,
)
from .config import ConnectionConfig, SafetyConfig, load_connection_map, load_safety_config
from .command_analysis import analyze_command as analyze_command_impl, parse_command

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

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


@dataclass
class RedisState:
    default_connection: ConnectionConfig
    connection_map: dict[str, ConnectionConfig]
    default_profile_name: str
    safety: SafetyConfig


@asynccontextmanager
async def lifespan(app: FastMCP) -> AsyncIterator[RedisState]:
    state = RedisState(
        default_connection=DEFAULT_CONN,
        connection_map=CONNECTION_MAP,
        default_profile_name=DEFAULT_PROFILE_NAME,
        safety=SAFETY,
    )
    print("Redis 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("Redis MCP Server 已关闭", file=sys.stderr)


mcp = FastMCP("Redis", lifespan=lifespan)


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

class ConnectionItem(BaseModel):
    profile: str = Field(description="profile 名")
    description: Optional[str] = Field(default=None, description="可选说明")
    mode: str = Field(description="standalone / cluster")


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


class KeyList(BaseModel):
    connection: str = Field(description="profile 名")
    pattern: str = Field(description="匹配模式")
    keys: list[str] = Field(description="key 列表")
    count: int = Field(description="数量")
    truncated: bool = Field(default=False, description="是否被 max_scan_keys 截断")


class InfoResult(BaseModel):
    connection: str = Field(description="profile 名")
    section: Optional[str] = Field(default=None, description="INFO 区段")
    data: dict[str, Any] = Field(description="INFO 返回的数据")


class CommandAnalysisResult(BaseModel):
    command: str = Field(description="命令名")
    args: list[str] = Field(description="参数列表")
    is_read: bool = Field(description="是否只读")
    is_write: bool = Field(description="是否写")
    is_dangerous: bool = Field(description="是否危险命令")
    need_confirm: bool = Field(description="是否需要 confirm 后执行")
    message: str = Field(description="说明")


class ExecuteCommandResult(BaseModel):
    command: str = Field(description="执行的命令")
    result: Any = Field(description="返回结果（可转为 JSON 兼容类型）")
    message: str = Field(default="", description="附加说明")


def _get_state() -> RedisState:
    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)


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

@mcp.tool()
def list_connections() -> ConnectionList:
    """列出所有可用 Redis 连接（profile、说明、单机/集群）。"""
    state = _get_state()
    items = [
        ConnectionItem(
            profile=k,
            description=getattr(c, "description", None),
            mode=c.mode,
        )
        for k, c in sorted(state.connection_map.items())
    ]
    return ConnectionList(connections=items, default=state.default_profile_name)


@mcp.tool()
def list_keys(
    pattern: Annotated[str, Field(description="key 匹配模式，如 * 或 user:*")] = "*",
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> KeyList:
    """使用 SCAN 列出 key，不超过配置的 max_scan_keys。"""
    state = _get_state()
    config = _conn(connection)
    profile_name = (connection or "").strip() or state.default_profile_name
    with connect(config, state.safety) as client:
        keys = _list_keys(client, pattern=pattern, max_keys=state.safety.max_scan_keys)
    truncated = len(keys) >= state.safety.max_scan_keys
    return KeyList(
        connection=profile_name,
        pattern=pattern,
        keys=keys,
        count=len(keys),
        truncated=truncated,
    )


@mcp.tool()
def info(
    section: Annotated[Optional[str], Field(description="INFO 区段，如 server/memory/clients，空为全部")] = None,
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> InfoResult:
    """返回 Redis INFO 信息。"""
    state = _get_state()
    config = _conn(connection)
    profile_name = (connection or "").strip() or state.default_profile_name
    with connect(config, state.safety) as client:
        data = info_sections(client, section)
    return InfoResult(connection=profile_name, section=section, data=data)


@mcp.tool()
def analyze_command(
    command_payload: Annotated[str, Field(description="一条 Redis 命令，如 GET foo 或 SET k v")],
) -> CommandAnalysisResult:
    """分析一条 Redis 命令（不执行）：是否只读/写/危险，是否需要 confirm。"""
    state = _get_state()
    a = analyze_command_impl(command_payload.strip(), state.safety.read_only)
    return CommandAnalysisResult(
        command=a.command,
        args=a.args,
        is_read=a.is_read,
        is_write=a.is_write,
        is_dangerous=a.is_dangerous,
        need_confirm=a.need_confirm,
        message=a.message,
    )


def _serialize_result(v: Any) -> Any:
    """将 Redis 返回转为 JSON 友好类型。"""
    if v is None or isinstance(v, (bool, int, float, str)):
        return v
    if isinstance(v, bytes):
        return v.decode("utf-8", errors="replace")
    if isinstance(v, (list, tuple)):
        return [_serialize_result(x) for x in v]
    if isinstance(v, dict):
        return {str(k): _serialize_result(x) for k, x in v.items()}
    return str(v)


@mcp.tool()
def execute_command(
    command_payload: Annotated[str, Field(description="单条 Redis 命令，由大模型产生，如 GET key 或 SET k v")],
    confirm: Annotated[Optional[bool], Field(description="写操作或需确认的命令必须为 true")] = None,
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> ExecuteCommandResult:
    """执行单条 Redis 命令。命令内容由大模型产生；MCP 只做安全检查与执行。写操作需 confirm=true。"""
    state = _get_state()
    payload = (command_payload or "").strip()
    if not payload:
        return ExecuteCommandResult(command="", result=None, message="空命令，不执行。")
    a = analyze_command_impl(payload, state.safety.read_only)
    if state.safety.read_only and (a.is_write or a.is_dangerous):
        return ExecuteCommandResult(command=a.command, result=None, message=a.message)
    if a.is_dangerous and a.command != "KEYS":
        return ExecuteCommandResult(command=a.command, result=None, message=a.message)
    if a.need_confirm and not confirm:
        return ExecuteCommandResult(command=a.command, result=None, message=a.message)
    cmd, args = parse_command(payload)
    config = _conn(connection)
    with connect(config, state.safety) as client:
        raw = _execute_command(client, cmd, *args)
    result = _serialize_result(raw)
    return ExecuteCommandResult(command=cmd, result=result, message="已执行。")


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

def _get_state_safe() -> RedisState | None:
    try:
        return mcp.get_context().request_context.lifespan_context
    except Exception:
        return None


@mcp.resource("redis://config")
def resource_config() -> str:
    """当前 Redis 配置与安全策略摘要（不暴露密码）。"""
    state = _get_state_safe()
    if not state:
        return "# 配置信息需在会话中获取。"
    s = state.safety
    lines = [
        "# Redis MCP 配置",
        f"- 默认 profile: {state.default_profile_name}",
        f"- 只读: {s.read_only}",
        f"- max_scan_keys: {s.max_scan_keys}",
        "",
        "## 连接",
    ]
    for k in sorted(state.connection_map):
        c = state.connection_map[k]
        lines.append(f"- {k}: {c.mode} {c.host}:{c.port} {c.description or ''}")
    return "\n".join(lines)


@mcp.resource("redis://keys/{pattern}")
def resource_keys(pattern: str) -> str:
    """指定 pattern 的 key 列表（SCAN，受 max_scan_keys 限制）。"""
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"
    config = state.default_connection
    with connect(config, state.safety) as client:
        keys = _list_keys(client, pattern=pattern or "*", max_keys=state.safety.max_scan_keys)
    if not keys:
        return "(无匹配 key)"
    return "\n".join(keys[:200]) + ("\n..." if len(keys) > 200 else "")


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

@mcp.prompt()
def redis_safe_usage_prompt() -> str:
    """Redis 安全使用习惯"""
    return """使用本 Redis MCP 时请遵守：
1. 先用 list_keys(pattern) 或 redis://keys/* 了解 key 空间，再生成命令。
2. 写命令（SET/DEL/INCR 等）执行前需传入 confirm=true。
3. 禁止通过 MCP 执行 FLUSHALL、FLUSHDB、CONFIG、DEBUG 等危险命令，请在 Redis 客户端中手动操作。
4. 尽量用 SCAN（list_keys）而非 KEYS，避免阻塞。
5. 单条命令由你生成，MCP 只做安全检查与执行。
"""


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

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


if __name__ == "__main__":
    main()
