"""
Jupyter Lab MCP Server - 重构版本
提供远程控制 Jupyter Lab 的 MCP 服务

提供的MCP工具：
1. connect_jupyter - 连接到Jupyter Lab服务器
2. disconnect_jupyter - 断开连接
3. get_connection_status - 获取连接状态
4. manage_notebook - 笔记本管理（list/create/get/update/delete/add_cell/update_cell/delete_cell/execute_cell）
5. manage_file - 文件管理（list/info/read/write/delete/download/upload/rename/mkdir）
6. manage_kernel - 内核管理（list/info，只读操作）
7. execute_code - 代码执行
"""

import json
from typing import Optional
from fastmcp import FastMCP

# 导入controller模块
from .controller import (
    session,
    notebook_manager,
    kernel_manager,
    file_manager,
    code_executor
)

# 创建 FastMCP 实例
mcp = FastMCP("Jupyter Lab MCP Server")


def create_response(success: bool = True, data: dict = None, 
                   message: str = None, error: str = None) -> str:
    """创建标准响应格式"""
    if success:
        response = {"success": True}
        if message:
            response["message"] = message
        if data:
            response.update(data)
    else:
        response = {
            "success": False,
            "error": error,
            "data": None
        }
    return json.dumps(response, ensure_ascii=False, indent=2)


# ============================================================================
# 会话管理工具
# ============================================================================

@mcp.tool()
def connect_jupyter(url: str = "http://localhost:8888", 
                   token: Optional[str] = None) -> str:
    """
    连接到 Jupyter Lab 服务器
    
    建立与 Jupyter Lab 服务器的连接，保存会话配置供后续操作使用。
    默认信任所有SSL证书，适合开发和内网环境。
    
    Args:
        url: Jupyter Lab 服务器地址
        token: 访问令牌
        
    Returns:
        JSON格式的连接结果
    """
    try:
        # 建立连接
        session.connect(url, token)
        
        # 测试连接
        from .controller.base import client
        server_info = client.request('/api/status')
        
        return create_response(
            success=True,
            message="成功连接到 Jupyter Lab 服务器",
            data={
                "server_info": server_info,
                "url": url,
                "has_token": token is not None
            }
        )
    except Exception as e:
        # 连接失败时清除会话
        session.disconnect()
        return create_response(success=False, error=str(e))


@mcp.tool()
def disconnect_jupyter() -> str:
    """
    断开 Jupyter Lab 连接
    
    清除保存的连接配置信息。
    
    Returns:
        JSON格式的断开连接结果
    """
    try:
        session.disconnect()
        return create_response(
            success=True,
            message="已断开 Jupyter Lab 连接"
        )
    except Exception as e:
        return create_response(success=False, error=str(e))


@mcp.tool()
def get_connection_status(include_server_info: bool = False) -> str:
    """
    获取当前连接状态
    
    显示当前保存的连接配置信息，可选择包含详细的服务器信息。
    
    Args:
        include_server_info: 是否包含详细的服务器信息
    
    Returns:
        JSON格式的连接状态
    """
    try:
        status = session.get_status()
        
        if include_server_info and status["connected"]:
            try:
                from .controller.base import client
                server_info = client.request('/api/status')
                status["server_info"] = server_info
                status["server_accessible"] = True
            except Exception as e:
                status["server_accessible"] = False
                status["server_error"] = str(e)
        
        return create_response(success=True, data=status)
    except Exception as e:
        return create_response(success=False, error=str(e))


# ============================================================================
# 笔记本管理工具
# ============================================================================

@mcp.tool()
def manage_notebook(action: str, path: Optional[str] = None,
                   kernel_name: Optional[str] = None,
                   content: Optional[str] = None,
                   source: Optional[str] = None,
                   cell_index: Optional[int] = None,
                   cell_type: str = "code",
                   position: Optional[int] = None) -> str:
    """
    笔记本管理工具（统一的笔记本操作接口）
    
    支持的操作：
    - "list": 列出所有笔记本
    - "create": 创建新笔记本（需要path，kernel_name可选，默认使用服务器默认内核）
    - "get": 获取笔记本内容（需要path）
    - "update": 更新整个笔记本（需要path和content，content为JSON字符串）
    - "delete": ⚠️ 删除笔记本（需要path）【危险操作，需要用户二次确认】
    - "add_cell": 添加单元格（需要path和source）
    - "update_cell": 更新单元格（需要path、cell_index和source）
    - "delete_cell": ⚠️ 删除单元格（需要path和cell_index）【需要用户二次确认】
    - "execute_cell": 执行指定单元格（需要path和cell_index）
    
    Args:
        action: 操作类型
        path: 笔记本路径
        kernel_name: 内核名称（create时可选，不提供则使用服务器默认内核）
        content: 完整的notebook内容JSON字符串（update时使用）
        source: 单元格内容（add_cell、update_cell时使用）
        cell_index: 单元格索引（update_cell、delete_cell、execute_cell时使用）
        cell_type: 单元格类型（add_cell时使用）
        position: 插入位置（add_cell时使用）
        
    Returns:
        JSON格式的操作结果
    """
    try:
        if not session.is_connected():
            return create_response(success=False, error="未连接到服务器，请先调用 connect_jupyter")
        
        if action == "list":
            notebooks = notebook_manager.list_all()
            return create_response(
                success=True,
                data={"notebooks": notebooks, "total": len(notebooks)}
            )
        
        elif action == "create":
            if not path:
                return create_response(success=False, error="创建笔记本需要提供path参数")
            result = notebook_manager.create(path, kernel_name)
            return create_response(
                success=True,
                message=f"笔记本 {path} 创建成功",
                data={"notebook": result}
            )
        
        elif action == "get":
            if not path:
                return create_response(success=False, error="获取笔记本需要提供path参数")
            result = notebook_manager.get(path)
            return create_response(success=True, data={"notebook": result})
        
        elif action == "update":
            if not path or not content:
                return create_response(success=False, error="更新笔记本需要提供path和content参数")
            try:
                notebook_content = json.loads(content)
            except json.JSONDecodeError as e:
                return create_response(success=False, error=f"content必须是有效的JSON字符串: {str(e)}")
            result = notebook_manager.update(path, notebook_content)
            return create_response(
                success=True,
                message=f"笔记本 {path} 更新成功",
                data={"notebook": result}
            )
        
        elif action == "delete":
            if not path:
                return create_response(success=False, error="删除笔记本需要提供path参数")
            result = notebook_manager.delete(path)
            return create_response(
                success=True,
                message=f"笔记本 {path} 删除成功",
                data=result
            )
        
        elif action == "add_cell":
            if not path or not source:
                return create_response(success=False, error="添加单元格需要提供path和source参数")
            result = notebook_manager.add_cell(path, source, cell_type, position)
            return create_response(
                success=True,
                message=f"成功向笔记本 {path} 添加{cell_type}单元格",
                data={"notebook": result}
            )
        
        elif action == "update_cell":
            if not path or cell_index is None or not source:
                return create_response(success=False, error="更新单元格需要提供path、cell_index和source参数")
            result = notebook_manager.update_cell(path, cell_index, source)
            return create_response(
                success=True,
                message=f"成功更新笔记本 {path} 的第 {cell_index} 个单元格",
                data={"notebook": result}
            )
        
        elif action == "delete_cell":
            if not path or cell_index is None:
                return create_response(success=False, error="删除单元格需要提供path和cell_index参数")
            result = notebook_manager.delete_cell(path, cell_index)
            return create_response(
                success=True,
                message=f"成功删除笔记本 {path} 的第 {cell_index} 个单元格",
                data={"notebook": result}
            )
        
        elif action == "execute_cell":
            if not path or cell_index is None:
                return create_response(success=False, error="执行单元格需要提供path和cell_index参数")
            
            # 获取notebook内容
            notebook = notebook_manager.get(path)
            cells = notebook.get('content', {}).get('cells', [])
            
            if cell_index < 0 or cell_index >= len(cells):
                return create_response(
                    success=False, 
                    error=f"单元格索引 {cell_index} 超出范围（共 {len(cells)} 个单元格）"
                )
            
            cell = cells[cell_index]
            if cell.get('cell_type') != 'code':
                return create_response(
                    success=False,
                    error=f"第 {cell_index} 个单元格不是代码单元格，类型为: {cell.get('cell_type')}"
                )
            
            # 获取单元格代码
            cell_source = cell.get('source', [])
            if isinstance(cell_source, list):
                code = ''.join(cell_source)
            else:
                code = cell_source
            
            # 执行代码
            result = code_executor.execute(path, code)
            return create_response(
                success=True,
                message=f"成功执行笔记本 {path} 的第 {cell_index} 个单元格",
                data={
                    "cell_index": cell_index,
                    "code": code,
                    **result
                }
            )
        
        else:
            return create_response(
                success=False,
                error=f"不支持的操作: {action}"
            )
    
    except Exception as e:
        return create_response(success=False, error=str(e))


# ============================================================================
# 文件管理工具
# ============================================================================

@mcp.tool()
def manage_file(action: str, path: str, 
               local_path: Optional[str] = None,
               content: Optional[str] = None,
               new_path: Optional[str] = None) -> str:
    """
    文件管理工具（统一的文件操作接口）
    
    支持的操作：
    - "list": 列出目录内容（path为目录路径）
    - "info": 获取文件信息（path为文件路径）
    - "read": 读取文件内容（path为文件路径）
    - "write": 写入文件内容（需要path和content）
    - "delete": ⚠️ 删除文件或目录（path为文件/目录路径）【危险操作，需要用户二次确认，删除后无法恢复】
    - "download": 下载文件到本地（path为远程路径，local_path为本地路径）
    - "upload": 上传本地文件（path为本地路径，new_path为远程路径）
    - "rename": 重命名文件（path为原路径，new_path为新路径）
    - "mkdir": 创建目录（path为目录路径）
    
    Args:
        action: 操作类型
        path: 文件/目录路径
        local_path: 本地路径（download时使用）
        content: 文件内容（write时使用）
        new_path: 新路径（upload、rename时使用）
        
    Returns:
        JSON格式的操作结果
    """
    try:
        if not session.is_connected():
            return create_response(success=False, error="未连接到服务器，请先调用 connect_jupyter")
        
        if action == "list":
            items = file_manager.list_directory(path)
            return create_response(
                success=True,
                data={"items": items, "total": len(items), "path": path}
            )
        
        elif action == "info":
            info = file_manager.get_info(path)
            return create_response(success=True, data={"file_info": info})
        
        elif action == "read":
            content_data = file_manager.read(path)
            return create_response(success=True, data={"file": content_data})
        
        elif action == "write":
            if not content:
                return create_response(success=False, error="写入文件需要提供content参数")
            result = file_manager.write(path, content)
            return create_response(
                success=True,
                message=f"文件 {path} 写入成功",
                data={"file": result}
            )
        
        elif action == "delete":
            result = file_manager.delete(path)
            return create_response(
                success=True,
                message=f"文件 {path} 删除成功",
                data=result
            )
        
        elif action == "download":
            result = file_manager.download(path, local_path)
            return create_response(
                success=True,
                message=f"文件下载成功",
                data=result
            )
        
        elif action == "upload":
            if not new_path:
                return create_response(success=False, error="上传文件需要提供new_path参数（远程路径）")
            result = file_manager.upload(path, new_path)
            return create_response(
                success=True,
                message=f"文件上传成功",
                data={"file": result}
            )
        
        elif action == "rename":
            if not new_path:
                return create_response(success=False, error="重命名文件需要提供new_path参数")
            result = file_manager.rename(path, new_path)
            return create_response(
                success=True,
                message=f"文件重命名成功",
                data={"file": result}
            )
        
        elif action == "mkdir":
            result = file_manager.create_directory(path)
            return create_response(
                success=True,
                message=f"目录 {path} 创建成功",
                data={"directory": result}
            )
        
        else:
            return create_response(
                success=False,
                error=f"不支持的操作: {action}"
            )
    
    except Exception as e:
        return create_response(success=False, error=str(e))


# ============================================================================
# 内核管理工具
# ============================================================================

@mcp.tool()
def manage_kernel(action: str, kernel_id: Optional[str] = None) -> str:
    """
    内核管理工具（统一的内核操作接口）
    
    支持的操作：
    - "list": 列出所有内核（包括运行中的内核和可用的内核规格）
    - "info": 获取内核信息（需要kernel_id）
    
    注意：为避免干扰服务器运行，不提供重启、中断、关闭等操作
    
    Args:
        action: 操作类型
        kernel_id: 内核ID（info操作需要）
        
    Returns:
        JSON格式的操作结果
    """
    try:
        if not session.is_connected():
            return create_response(success=False, error="未连接到服务器，请先调用 connect_jupyter")
        
        if action == "list":
            kernel_info = kernel_manager.list_all()
            return create_response(
                success=True,
                data=kernel_info
            )
        
        elif action == "info":
            if not kernel_id:
                return create_response(success=False, error="获取内核信息需要提供kernel_id参数")
            info = kernel_manager.get_info(kernel_id)
            return create_response(success=True, data={"kernel": info})
        
        else:
            return create_response(
                success=False,
                error=f"不支持的操作: {action}。仅支持 list 和 info"
            )
    
    except Exception as e:
        return create_response(success=False, error=str(e))


# ============================================================================
# 代码执行工具
# ============================================================================

@mcp.tool()
def execute_code(notebook_path: str, code: str) -> str:
    """
    执行代码
    
    在指定笔记本中执行Python代码。
    
    Args:
        notebook_path: 笔记本路径
        code: 要执行的代码
        
    Returns:
        JSON格式的执行结果
    """
    try:
        if not session.is_connected():
            return create_response(success=False, error="未连接到服务器，请先调用 connect_jupyter")
        
        result = code_executor.execute(notebook_path, code)
        return create_response(
            success=True,
            message="代码已提交执行",
            data=result
        )
    
    except Exception as e:
        return create_response(success=False, error=str(e))


# ============================================================================
# 主函数
# ============================================================================

if __name__ == "__main__":
    mcp.run()


def main():
    """命令行入口点函数"""
    mcp.run()
