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

敏感信息（hosts、用户名、密码）通过环境变量名引用，不暴露实际值。
首次使用请先运行 setup 工具进行配置。
"""

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 (
    AppConfig,
    ConnectionProfile,
    SafetyConfig,
    get_connection_values,
    load_config,
    need_setup,
)
from .request_analysis import analyze_request as analyze_request_impl
from .tools import do_list_connections, do_get_connection, run_setup, do_check_setup_status


# 加载配置
APP_CONFIG = load_config()
SAFETY = APP_CONFIG.safety


def _build_connection_config(profile: ConnectionProfile) -> Any:
    """从 ConnectionProfile 构建兼容的 ConnectionConfig"""
    from .client import ConnectionConfig

    values = get_connection_values(profile)

    # 解析 hosts URL
    from urllib.parse import urlparse
    hosts_url = values.get("hosts", "").strip()
    if hosts_url:
        if "://" not in hosts_url:
            hosts_url = "http://" + hosts_url
        p = urlparse(hosts_url)
        host = p.hostname or p.netloc or ""
        port = p.port if p.port else 9200
        use_ssl = p.scheme.lower() == "https" if p.scheme else False
    else:
        host = ""
        port = 9200
        use_ssl = False

    return ConnectionConfig(
        host=host,
        port=port,
        user=values.get("username", ""),
        password=values.get("password", ""),
        use_ssl=use_ssl,
        description=profile.description,
        skip_product_check=profile.skip_product_check,
    )


@dataclass
class ESState:
    app_config: AppConfig
    connection_configs: dict[str, Any]  # name -> ConnectionConfig
    safety: SafetyConfig


@asynccontextmanager
async def lifespan(app: FastMCP) -> AsyncIterator[ESState]:
    """生命周期管理 - 启动时检查配置"""
    # 构建连接配置
    connection_configs = {}
    for profile in APP_CONFIG.profiles:
        try:
            connection_configs[profile.name] = _build_connection_config(profile)
        except Exception as e:
            print(f"Warning: Failed to build config for {profile.name}: {e}", file=sys.stderr)

    state = ESState(
        app_config=APP_CONFIG,
        connection_configs=connection_configs,
        safety=SAFETY,
    )

    # 检查是否需要引导配置
    if need_setup():
        print("=" * 60, file=sys.stderr)
        print("Elasticsearch MCP 未配置或需要重新配置", file=sys.stderr)
        print("请运行 setup 工具进行配置", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
    else:
        print("Elasticsearch MCP Server 已启动", file=sys.stderr)
        if APP_CONFIG.profiles:
            default = APP_CONFIG.default_profile()
            print(f"  默认 profile: {default.name if default else '无'}", file=sys.stderr)

    yield state

    print("Elasticsearch MCP Server 已关闭", file=sys.stderr)


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


# ============ Pydantic Models ============

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


class ConnectionList(BaseModel):
    connections: list[ConnectionItem]
    count: int
    default: Optional[str] = None


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: Any
    message: str = ""


# ============ 内部函数 ============

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


def _conn(profile_name: Optional[str]) -> Any:
    """根据 profile 名称获取连接配置"""
    state = _get_state()
    if profile_name:
        config = state.connection_configs.get(profile_name)
        if not config:
            # 尝试构建
            p = state.app_config.get_profile(profile_name)
            if p:
                config = _build_connection_config(p)
        return config or state.connection_configs.get(state.app_config.default_profile().name if state.app_config.default_profile() else "")
    return state.connection_configs.get(state.app_config.default_profile().name if state.app_config.default_profile() else "")


# ============ 工具 ============

@mcp.tool()
def list_connections() -> dict:
    """
    列出所有已配置的 ES 连接。

    注意：敏感信息（hosts、用户名、密码）仅显示环境变量名，不显示实际值。
    """
    state = _get_state()

    connections = []
    for profile in state.app_config.profiles:
        connections.append({
            "profile": profile.name,
            "env": profile.env,
            "description": profile.description or "无",
            "has_auth": profile.has_auth(),
            "skip_product_check": profile.skip_product_check,
            "env_vars": {
                k: v for k, v in {
                    "hosts": profile.env_vars.hosts,
                    "username": profile.env_vars.username,
                    "password": profile.env_vars.password,
                }.items() if v is not None
            },
        })

    default = state.app_config.default_profile()

    return {
        "connections": connections,
        "count": len(connections),
        "default": default.name if default else None,
    }


@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.app_config.default_profile().name if state.app_config.default_profile() else "")

    if not config:
        return IndexList(connection=profile_name, pattern=pattern, indices=[], count=0, message="未找到连接配置，请先运行 setup")

    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.app_config.default_profile().name if state.app_config.default_profile() else "")

    if not config:
        return MappingResult(connection=profile_name, index=index, mapping={}, message="未找到连接配置，请先运行 setup")

    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)
    if not config:
        return ExecuteRequestResult(method=method, path=path_norm, result={}, message="未找到连接配置，请先运行 setup")

    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="已执行。")


@mcp.tool()
def setup(
    description: Annotated[str | None, "连接描述，如：测试环境ES"] = None,
    env_name: Annotated[str, "环境名称: test/uat/prod"] = "test",
    hosts_env_var: Annotated[str, "hosts 环境变量名（运行时读取）"] = "",
    username_env_var: Annotated[str | None, "username 环境变量名（可为空）"] = None,
    password_env_var: Annotated[str, "password 环境变量名（运行时读取）"] = "",
) -> dict:
    """
    配置 ES 连接。

    敏感信息（hosts、用户名、密码）通过环境变量名引用，不暴露实际值。
    配置完成后会输出环境变量设置命令，请在您的 shell 配置文件中填入实际值。
    """
    return run_setup(
        description=description,
        env_name=env_name,
        hosts_env_var=hosts_env_var or None,
        username_env_var=username_env_var,
        password_env_var=password_env_var or None,
    )


@mcp.tool()
def check_setup_status() -> dict:
    """
    检查当前配置状态。
    """
    return do_check_setup_status()


# ============ 资源 ============

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.app_config.default_profile().name if state.app_config.default_profile() else '无'}",
        f"- 只读: {state.safety.read_only}",
        f"- max_search_size: {state.safety.max_search_size}",
        "",
        "# 已配置的连接：",
    ]

    for profile in state.app_config.profiles:
        auth_info = "有认证" if profile.has_auth() else "无认证"
        lines.append(f"- {profile.name} ({profile.env}): {profile.description or '无'} - {auth_info}")
        if profile.env_vars.hosts:
            lines.append(f"  hosts: {profile.env_vars.hosts} (环境变量)")

    return "\n".join(lines)


@mcp.resource("es://indices/{pattern}")
def resource_indices(pattern: str) -> str:
    """获取索引列表"""
    state = _get_state_safe()
    if not state:
        return "(需在会话中加载)"

    default = state.app_config.default_profile()
    if not default:
        return "(无默认连接，请先运行 setup)"

    config = state.connection_configs.get(default.name)
    if not config:
        return "(连接配置无效)"

    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_connections 查看可用连接，再用 list_indices、get_mapping 了解索引与字段，再构造请求。
2. 写请求（index/update/delete）执行前需确认。
3. 敏感信息（hosts、用户名、密码）通过环境变量引用，不直接暴露。
4. 每次执行请求前会进行安全检查，危险操作会被拦截。

如需配置新连接，请运行 setup 工具。
"""


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


if __name__ == "__main__":
    main()
