#!/usr/bin/env python3
"""
Elasticsearch MCP Server - 索引发现、mapping、安全执行（请求由大模型产生）。
"""

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_request as _execute_request,
    get_connection_config,
    get_mapping as _get_mapping,
    list_indices as _list_indices,
)
from .config import ConnectionConfig, SafetyConfig, load_connection_map, load_safety_config
from .request_analysis import analyze_request as analyze_request_impl

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


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


@asynccontextmanager
async def lifespan(app: FastMCP) -> AsyncIterator[ESState]:
    state = ESState(
        default_connection=DEFAULT_CONN,
        connection_map=CONNECTION_MAP,
        default_profile_name=DEFAULT_PROFILE_NAME,
        safety=SAFETY,
    )
    print("Elasticsearch MCP Server 已启动", file=sys.stderr)
    print(f"  默认 profile: {DEFAULT_PROFILE_NAME}", file=sys.stderr)
    yield state
    print("Elasticsearch MCP Server 已关闭", file=sys.stderr)


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


class ConnectionItem(BaseModel):
    profile: str
    description: Optional[str] = None
    host: str = ""
    port: int = 9200


class ConnectionList(BaseModel):
    connections: list[ConnectionItem]
    default: str


class IndexList(BaseModel):
    connection: str
    pattern: str
    indices: list[dict]
    count: int


class MappingResult(BaseModel):
    connection: str
    index: str
    mapping: dict[str, Any]


class RequestAnalysisResult(BaseModel):
    method: str
    path: str
    is_read: bool
    is_write: bool
    is_dangerous: bool
    need_confirm: bool
    message: str


class ExecuteRequestResult(BaseModel):
    method: str
    path: str
    result: dict[str, Any]
    message: str = ""


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


@mcp.tool()
def list_connections() -> ConnectionList:
    """列出所有可用 ES 连接。"""
    state = _get_state()
    items = [
        ConnectionItem(
            profile=k,
            description=c.description,
            host=c.host,
            port=c.port,
        )
        for k, c in sorted(state.connection_map.items())
    ]
    return ConnectionList(connections=items, default=state.default_profile_name)


@mcp.tool()
def list_indices(
    pattern: Annotated[str, Field(description="索引名匹配，如 * 或 log-*")] = "*",
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> IndexList:
    """列出索引（支持 pattern）。"""
    state = _get_state()
    config = _conn(connection)
    profile_name = (connection or "").strip() or state.default_profile_name
    with connect(config, state.safety) as client:
        indices = _list_indices(client, pattern=pattern)
    return IndexList(connection=profile_name, pattern=pattern, indices=indices, count=len(indices))


@mcp.tool()
def get_mapping(
    index: Annotated[str, Field(description="索引名")],
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> MappingResult:
    """获取索引 mapping。"""
    state = _get_state()
    config = _conn(connection)
    profile_name = (connection or "").strip() or state.default_profile_name
    with connect(config, state.safety) as client:
        mapping = _get_mapping(client, index)
    return MappingResult(connection=profile_name, index=index, mapping=mapping)


@mcp.tool()
def analyze_request(
    method: Annotated[str, Field(description="HTTP 方法，如 GET、POST")],
    path: Annotated[str, Field(description="路径，如 /my_index/_search")],
    body: Annotated[Optional[dict], Field(description="请求体，搜索 DSL 等")] = None,
) -> RequestAnalysisResult:
    """分析一条 ES 请求（不执行）：只读/写/危险，是否需 confirm。"""
    state = _get_state()
    a = analyze_request_impl(method, path, body, state.safety.read_only)
    return RequestAnalysisResult(
        method=a.method,
        path=a.path,
        is_read=a.is_read,
        is_write=a.is_write,
        is_dangerous=a.is_dangerous,
        need_confirm=a.need_confirm,
        message=a.message,
    )


@mcp.tool()
def execute_request(
    method: Annotated[str, Field(description="HTTP 方法")],
    path: Annotated[str, Field(description="路径，如 /my_index/_search")],
    body: Annotated[Optional[dict], Field(description="请求体")] = None,
    confirm: Annotated[Optional[bool], Field(description="写操作必须为 true")] = None,
    connection: Annotated[Optional[str], Field(description="profile 名")] = None,
) -> ExecuteRequestResult:
    """执行单条 ES 请求。请求由大模型产生；MCP 只做安全检查与执行。写操作需 confirm=true。"""
    state = _get_state()
    path_norm = (path or "").strip().rstrip("/") or "/"
    if not path_norm.startswith("/"):
        path_norm = "/" + path_norm
    a = analyze_request_impl(method, path_norm, body, state.safety.read_only)
    if state.safety.read_only and (a.is_write or a.is_dangerous):
        return ExecuteRequestResult(method=method, path=path_norm, result={}, message=a.message)
    if a.is_dangerous:
        return ExecuteRequestResult(method=method, path=path_norm, result={}, message=a.message)
    if a.need_confirm and not confirm:
        return ExecuteRequestResult(method=method, path=path_norm, result={}, message=a.message)
    if a.is_read and body and isinstance(body, dict) and "_search" in path_norm:
        size = body.get("size")
        if size is not None and int(size) > state.safety.max_search_size:
            body = {**body, "size": state.safety.max_search_size}
    config = _conn(connection)
    with connect(config, state.safety) as client:
        result = _execute_request(client, method.upper(), path_norm, body)
    return ExecuteRequestResult(method=method, path=path_norm, result=result, message="已执行。")


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


@mcp.resource("es://config")
def resource_config() -> str:
    state = _get_state_safe()
    if not state:
        return "# 需在会话中加载"
    lines = ["# Elasticsearch MCP 配置", f"- 默认 profile: {state.default_profile_name}", f"- 只读: {state.safety.read_only}", f"- max_search_size: {state.safety.max_search_size}", ""]
    for k in sorted(state.connection_map):
        c = state.connection_map[k]
        lines.append(f"- {k}: {c.host}:{c.port} {c.description or ''}")
    return "\n".join(lines)


@mcp.resource("es://indices/{pattern}")
def resource_indices(pattern: str) -> str:
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"
    config = state.default_connection
    with connect(config, state.safety) as client:
        indices = _list_indices(client, pattern=pattern or "*")
    if not indices:
        return "(无匹配索引)"
    return "\n".join([f"- {x.get('name', '')} docs={x.get('docs_count', '')}" for x in indices[:100]])


@mcp.prompt()
def es_safe_usage_prompt() -> str:
    return """使用本 Elasticsearch MCP 时请遵守：
1. 先用 list_indices、get_mapping 了解索引与字段，再构造请求。
2. 写请求（index/update/delete）执行前需传入 confirm=true。
3. 禁止删除整个索引或操作 _cluster/_nodes，请在 Kibana/控制台操作。
4. 单条请求（method+path+body）由你产生，MCP 只做安全检查与执行。
"""


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


if __name__ == "__main__":
    main()
