# topological_analysis_tools.py

import posixpath
import re
from typing import Any, Literal
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator

import asyncssh
from pydantic import BaseModel, Field

# 从 MCP SDK 导入核心组件 [1]
from mcp.server.fastmcp import FastMCP, Context
from mcp.server.session import ServerSession

# 从 lifespan.py 导入共享的配置和数据类 [3]
# 注意: 我们不再需要导入 system_lifespan
from lifespan import (
    SharedAppContext,
    REMOTE_HOST,
    REMOTE_USER,
    PRIVATE_KEY_PATH,
    INITIAL_WORKING_DIRECTORY,
)

# ==============================================================================
# 1. 辅助函数 (为实现独立性而复制)
# ==============================================================================

def shell_quote(arg: str) -> str:
    """安全地将字符串转义为单个 shell 参数 [5]。"""
    return "'" + arg.replace("'", "'\"'\"'") + "'"

def _safe_join(sandbox_root: str, relative_path: str) -> str:
    """安全地将用户提供的相对路径连接到沙箱根目录 [5]。"""
    rel = (relative_path or ".").strip()
    if rel.startswith("/"):
        rel = rel.lstrip("/")
    combined = posixpath.normpath(posixpath.join(sandbox_root, rel))
    root_norm = sandbox_root.rstrip("/")
    if combined != root_norm and not combined.startswith(root_norm + "/"):
        raise ValueError("路径越界: 仅允许访问沙箱目录内部")
    if ".." in combined.split("/"):
        raise ValueError("非法路径: 不允许使用 '..' 跨目录")
    return combined

# ==============================================================================
# 2. 独立的生命周期管理器
# ==============================================================================

@asynccontextmanager
async def analysis_lifespan(_server: FastMCP) -> AsyncIterator[SharedAppContext]:
    """
    为拓扑分析工具独立管理 SSH 连接的生命周期。
    """
    conn: asyncssh.SSHClientConnection | None = None
    print("分析工具: 正在建立独立的 SSH 连接...")
    try:
        conn = await asyncssh.connect(
            REMOTE_HOST,
            username=REMOTE_USER,
            client_keys=[PRIVATE_KEY_PATH],
        )
        print(f"分析工具: 独立的 SSH 连接到 {REMOTE_HOST} 成功!")
        yield SharedAppContext(ssh_connection=conn, sandbox_path=INITIAL_WORKING_DIRECTORY)
    finally:
        if conn:
            conn.close()
            await conn.wait_closed()
            print("分析工具: 独立的 SSH 连接已关闭。")


# ==============================================================================
# 3. 数据模型与解析函数 (与之前相同)
# ==============================================================================
# (为简洁起见，此处省略了 Pydantic 模型和 _parse_analysis_output 函数的定义，
#  请将上一版本中的这部分代码粘贴到这里)
VALID_CONFIGS = Literal["O.yaml", "S.yaml", "Cl.yaml", "Br.yaml"]
class SubmitTopologicalAnalysisParams(BaseModel):
    cif_file_path: str = Field(description="远程服务器上 CIF 文件的相对路径。")
    config_files: list[VALID_CONFIGS] = Field(description="选择一个或多个用于分析的 YAML 配置文件。")
    output_file_path: str = Field(description="用于保存计算结果的输出文件相对路径。")
class AnalysisResult(BaseModel):
    percolation_diameter_a: float | None = Field(None, description="渗透直径 (Percolation diameter), 单位 Å。")
    connectivity_distance: float | None = Field(None, description="连通距离 (the minium of d), 单位 Å。")
    max_node_length_a: float | None = Field(None, description="最大节点长度 (Maximum node length), 单位 Å。")
    min_cation_distance_a: float | None = Field(None, description="到阳离子的最小距离, 单位 Å。")
    total_time_s: float | None = Field(None, description="总计算耗时, 单位秒。")
    long_node_count: int | None = Field(None, description="长节点数量 (Long node number)。")
    short_node_count: int | None = Field(None, description="短节点数量 (Short node number)。")
    warnings: list[str] = Field(default_factory=list, description="解析过程中遇到的警告或缺失的关键值信息。")
def _parse_analysis_output(log_content: str) -> AnalysisResult:
    def find_float(pattern: str) -> float | None:
        match = re.search(pattern, log_content)
        return float(match.group(1)) if match else None
    def find_int(pattern: str) -> int | None:
        match = re.search(pattern, log_content)
        return int(match.group(1)) if match else None
    data = { "percolation_diameter_a": find_float(r"Percolation diameter \(A\):\s*([\d\.]+)"), "max_node_length_a": find_float(r"Maximum node length detected:\s*([\d\.]+)\s*A"), "min_cation_distance_a": find_float(r"minimum distance to cations is\s*([\d\.]+)\s*A"), "total_time_s": find_float(r"Total used time:\s*([\d\.]+)"), "long_node_count": find_int(r"Long node number:\s*(\d+)"), "short_node_count": find_int(r"Short node number:\s*(\d+)"), }
    conn_dist_match = re.search(r"the minium of d\s*([\d\.]+)", log_content, re.MULTILINE)
    data["connectivity_distance"] = float(conn_dist_match.group(1)) if conn_dist_match else None
    warnings = [k + " 未找到" for k, v in data.items() if v is None and k not in ["min_cation_distance_a", "long_node_count", "short_node_count"]]
    return AnalysisResult(**data, warnings=warnings)


# ==============================================================================
# 4. MCP 服务器工厂函数
# ==============================================================================

def create_topological_analysis_mcp() -> FastMCP:
    """
    供 Starlette 挂载的工厂函数。
    这个函数现在创建并使用自己独立的 lifespan 管理器。
    """
    analysis_mcp = FastMCP(
        name="Topological Analysis Tools",
        instructions="用于在远程服务器上提交和分析拓扑计算任务的工具集。",
        streamable_http_path="/",
        lifespan=analysis_lifespan, # 关键: 使用本文件内定义的独立 lifespan
    )

    @analysis_mcp.tool()
    async def submit_topological_analysis(
        params: SubmitTopologicalAnalysisParams,
        ctx: Context[ServerSession, SharedAppContext],
    ) -> dict[str, Any]:
        """【步骤1/2】异步提交拓扑分析任务。"""
        # ... (工具的实现逻辑与之前完全相同)
        try:
            app_ctx = ctx.request_context.lifespan_context
            conn = app_ctx.ssh_connection
            sandbox_root = app_ctx.sandbox_path
            home_dir = f"/cluster/home/{REMOTE_USER}"
            tool_dir = f"{home_dir}/tool/Topological_analyse"
            cif_abs_path = _safe_join(sandbox_root, params.cif_file_path)
            output_abs_path = _safe_join(sandbox_root, params.output_file_path)
            config_args = " ".join([shell_quote(f"{tool_dir}/{cfg}") for cfg in params.config_files])
            config_flags = f"-i {config_args}"
            command = f"""
            nohup sh -c '
                source {shell_quote(f"{home_dir}/.bashrc")} && \\
                conda activate {shell_quote(f"{home_dir}/anaconda3/envs/zeo")} && \\
                python {shell_quote(f"{tool_dir}/analyze_voronoi_nodes.py")} \\
                       {shell_quote(cif_abs_path)} \\
                       {config_flags}
            ' > {shell_quote(output_abs_path)} 2>&1 &
            """.strip()
            await ctx.info("提交后台拓扑分析任务...")
            await conn.run(command, check=True)
            return {"success": True, "message": "拓扑分析任务已成功提交。", "output_file": params.output_file_path}
        except Exception as e:
            msg = f"提交后台任务失败: {e}"
            await ctx.error(msg)
            return {"success": False, "error": msg}

    @analysis_mcp.tool()
    async def analyze_topological_results(
        output_file_path: str,
        ctx: Context[ServerSession, SharedAppContext],
    ) -> AnalysisResult:
        """【步骤2/2】读取并分析拓扑分析任务的输出文件。"""
        # ... (工具的实现逻辑与之前完全相同)
        try:
            await ctx.info(f"开始分析结果文件: {output_file_path}")
            app_ctx = ctx.request_context.lifespan_context
            conn = app_ctx.ssh_connection
            sandbox_root = app_ctx.sandbox_path
            target_path = _safe_join(sandbox_root, output_file_path)
            async with conn.start_sftp_client() as sftp:
                async with sftp.open(target_path, "r") as f:
                    log_content = await f.read()
            if not log_content:
                raise ValueError("结果文件为空。")
            analysis_data = _parse_analysis_output(log_content)
            return analysis_data
        except FileNotFoundError:
            msg = f"分析失败: 结果文件 '{output_file_path}' 不存在。"
            await ctx.error(msg)
            return AnalysisResult(warnings=[msg])
        except Exception as e:
            msg = f"分析结果时发生错误: {e}"
            await ctx.error(msg)
            return AnalysisResult(warnings=[msg])

    return analysis_mcp