"""
WebSocket Client Module

Establishes WebSocket connection with server, receives commands and returns execution results
"""

import asyncio
import json
import logging
import os
import time
import uuid
from typing import Optional

import aiohttp

from ..operations import document
from ..core.executor import CommandExecutor
from ..core.openclaw import OpenClawChecker
from ..core.config import get_config
from ..mcp.manager import get_mcp_manager
from ..security.whitelist import get_whitelist_file_path, is_command_allowed

# Configure logging
logger = logging.getLogger(__name__)


class WebSocketClient:
    """WebSocket client"""

    def __init__(
        self,
        server_url: str,
        client_id: str,
        verification_code: str,
        workspace_id: str,
        reconnect_interval: float = 5.0,
    ):
        """
        初始化 WebSocket client

        Args:
            server_url: Server URL
            client_id: Client unique identifier
            verification_code: Verification code
            workspace_id: Workspace ID
            reconnect_interval: Reconnect interval (seconds)
        """
        self.server_url = server_url
        self.client_id = client_id
        self.verification_code = verification_code
        self.workspace_id = workspace_id
        self.reconnect_interval = reconnect_interval
        self.mcp_manager = get_mcp_manager()
        self.running = False
        self.ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self.session: Optional[aiohttp.ClientSession] = None
        self.heartbeat_task: Optional[asyncio.Task] = None
        self.heartbeat_interval = 30.0  # 30seconds，Ensure enough heartbeats before server timeout
        # OpenClaw check related
        self.openclaw_checker = OpenClawChecker()
        self.openclaw_check_enabled = get_config().get_openclaw_check_enabled()
        self.openclaw_check_task: Optional[asyncio.Task] = None
        self.openclaw_check_interval = 60.0  # OpenClaw check interval (seconds)
        self.last_openclaw_status = None  # Track last status to avoid duplicate notifications

    def _get_ws_url(self) -> str:
        """Get WebSocket URL"""
        # Convert http/https to ws/wss
        ws_url = self.server_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = ws_url.rstrip("/")

        # Add client authentication parameters
        return f"{ws_url}/api/v1/bridge/ws?client_id={self.client_id}&workspace_id={self.workspace_id}&verification_code={self.verification_code}"

    def _log_operation(
        self,
        session_id: str,
        operation_type: str,
        operation: str,
        arguments: any = None,
        result: any = None,
        success: bool = True,
        elapsed: float = 0.0,
        error: str = None,
    ):
        """
        Log operation to terminal.log

        Args:
            session_id: Session ID
            operation_type: Operation type（MCP_TOOL, DOC_READ, DOC_WRITE, DOC_EDIT）
            operation: Operation name
            arguments: Arguments
            result: Result
            success: Whether successful
            elapsed: Execution time
            error: Error message
        """
        from datetime import datetime

        try:
            # Create session directory
            home_dir = os.path.expanduser("~")
            session_dir = os.path.join(home_dir, ".castrel", session_id)
            os.makedirs(session_dir, exist_ok=True)
            log_file = os.path.join(session_dir, "terminal.log")

            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            log_entry = f"""[{timestamp}] {operation_type}: {operation}
  SUCCESS: {success}
  DURATION: {elapsed:.2f}s
"""

            if arguments:
                import json

                log_entry += f"  ARGUMENTS:\n    {json.dumps(arguments, ensure_ascii=False, indent=2).replace(chr(10), chr(10) + '    ')}\n"

            if result:
                import json

                result_str = json.dumps(result, ensure_ascii=False, indent=2)[:500]  # Limit length
                log_entry += f"  RESULT:\n    {result_str.replace(chr(10), chr(10) + '    ')}\n"

            if error:
                log_entry += f"  ERROR:\n    {error}\n"

            log_entry += "---\n\n"

            # Append to log file
            with open(log_file, "a", encoding="utf-8") as f:
                f.write(log_entry)

        except Exception as e:
            # Logging failure should not affect operation execution
            logger.warning(f"Failed to log operation: {e}")

    async def _handle_command(self, message: dict) -> Optional[dict]:
        """
        Handle server commands

        Args:
            message: Server message

        Returns:
            Optional[dict]: Response message, return None if no response needed None
        """
        message_id = message.get("id")
        message_type = message.get("type")
        timestamp = message.get("timestamp")

        logger.info(
            f"[CLIENT-MSG-RECV] Received message: message_id={message_id}, "
            f"message_type={message_type}, timestamp={timestamp}, client_id={self.client_id}"
        )

        if message_type == "connected":
            # Connection success message
            session_id = message.get("session_id", "")
            msg = message.get("message", "")
            logger.info(
                f"[CLIENT-CONNECTED] Connection established: session_id={session_id}, "
                f"message={msg}, client_id={self.client_id}"
            )
            # No response needed
            return None

        elif message_type == "local_tool_call":
            # Handle local command call
            data = message.get("data", {})
            command = data.get("command", "")
            args = data.get("args", [])
            cwd = data.get("cwd")
            timeout = data.get("timeout", 300)
            session_id = data.get("session_id", "")

            logger.info(
                f"[CLIENT-LOCAL-CALL] Local tool call received: message_id={message_id}, "
                f"command={command}, args={args}, cwd={cwd}, session_id={session_id}, timeout={timeout}s, "
                f"client_id={self.client_id}"
            )

            return await self._execute_local_command(
                message_id=message_id,
                command=command,
                args=args,
                cwd=cwd,
                session_id=session_id,
                timeout=timeout,
            )

        elif message_type == "mcp_tool_call":
            # Handle MCP tool call
            data = message.get("data", {})
            server_name = data.get("server_name", "")
            tool_name = data.get("tool_name", "")
            arguments = data.get("arguments", {})
            session_id = data.get("session_id", "")

            logger.info(
                f"[CLIENT-MCP-CALL] MCP tool call received: message_id={message_id}, "
                f"server={server_name}, tool={tool_name}, session_id={session_id}, arguments={arguments}, "
                f"client_id={self.client_id}"
            )

            return await self._execute_mcp_tool(
                message_id=message_id,
                server_name=server_name,
                tool_name=tool_name,
                session_id=session_id,
                arguments=arguments,
            )

        elif message_type == "doc_read_call":
            # Handle document read call
            data = message.get("data", {})
            file_path = data.get("file_path", "")
            encoding = data.get("encoding")
            session_id = data.get("session_id", "")

            logger.info(
                f"[CLIENT-DOC-READ-CALL] Doc read call received: message_id={message_id}, "
                f"file_path={file_path}, session_id={session_id}, encoding={encoding}, client_id={self.client_id}"
            )

            return await self._execute_doc_read(
                message_id=message_id,
                file_path=file_path,
                session_id=session_id,
                encoding=encoding,
            )

        elif message_type == "doc_write_call":
            # Handle document write call
            data = message.get("data", {})
            file_path = data.get("file_path", "")
            content = data.get("content", "")
            encoding = data.get("encoding", "utf-8")
            create_dirs = data.get("create_dirs", True)
            session_id = data.get("session_id", "")

            logger.info(
                f"[CLIENT-DOC-WRITE-CALL] Doc write call received: message_id={message_id}, "
                f"file_path={file_path}, content_len={len(content)}, session_id={session_id}, encoding={encoding}, "
                f"create_dirs={create_dirs}, client_id={self.client_id}"
            )

            return await self._execute_doc_write(
                message_id=message_id,
                file_path=file_path,
                session_id=session_id,
                content=content,
                encoding=encoding,
                create_dirs=create_dirs,
            )

        elif message_type == "doc_edit_call":
            # Handle document edit call
            data = message.get("data", {})
            file_path = data.get("file_path", "")
            operation = data.get("operation", "")
            new_content = data.get("new_content", "")
            old_content = data.get("old_content")
            encoding = data.get("encoding")
            session_id = data.get("session_id", "")

            logger.info(
                f"[CLIENT-DOC-EDIT-CALL] Doc edit call received: message_id={message_id}, "
                f"file_path={file_path}, operation={operation}, session_id={session_id}, encoding={encoding}, client_id={self.client_id}"
            )

            return await self._execute_doc_edit(
                message_id=message_id,
                file_path=file_path,
                session_id=session_id,
                operation=operation,
                new_content=new_content,
                old_content=old_content,
                encoding=encoding,
            )

        elif message_type == "ping":
            # Heartbeat from server, response needed
            logger.debug(f"[CLIENT-PING-RECV] Received ping: message_id={message_id}, client_id={self.client_id}")
            return {"id": message_id, "type": "pong"}

        elif message_type == "pong":
            # Server response to client heartbeat
            logger.debug(f"[CLIENT-PONG-RECV] Received pong: message_id={message_id}, client_id={self.client_id}")
            # No response needed
            return None

        else:
            # Unknown command type
            logger.warning(
                f"[CLIENT-MSG-UNKNOWN] Unknown message type: message_id={message_id}, "
                f"message_type={message_type}, client_id={self.client_id}"
            )
            return {
                "id": message_id,
                "type": "error",
                "error": f"Unknown message type: {message_type}",
            }

    async def _send_heartbeat(self):
        """Send heartbeat periodically"""
        logger.info(
            f"[CLIENT-HEARTBEAT-START] Heartbeat task started: interval={self.heartbeat_interval}s, "
            f"client_id={self.client_id}"
        )

        while self.running and self.ws and not self.ws.closed:
            try:
                # Send heartbeat message
                heartbeat_msg = {
                    "id": str(uuid.uuid4()),
                    "type": "ping",
                    "timestamp": int(time.time() * 1000),
                }
                logger.debug(
                    f"[CLIENT-HEARTBEAT-SEND] Sending heartbeat: message_id={heartbeat_msg['id']}, "
                    f"client_id={self.client_id}"
                )
                await self.ws.send_json(heartbeat_msg)
                logger.debug(
                    f"[CLIENT-HEARTBEAT-SENT] Heartbeat sent: message_id={heartbeat_msg['id']}, "
                    f"client_id={self.client_id}"
                )

                # Wait for next heartbeat
                await asyncio.sleep(self.heartbeat_interval)

            except Exception as e:
                logger.error(
                    f"[CLIENT-HEARTBEAT-ERROR] Failed to send heartbeat: error={e}, client_id={self.client_id}",
                    exc_info=True,
                )
                break

        logger.info(f"[CLIENT-HEARTBEAT-STOP] Heartbeat task stopped: client_id={self.client_id}")

    async def _perform_openclaw_check(self) -> dict:
        """
        Perform OpenClaw check
        
        Returns:
            dict: OpenClaw status with the following structure:
                {
                    "status": "healthy|warning|error",
                    "message": "Status description",
                    "details": {
                        # Check-specific parameters from OpenClawChecker
                    }
                }
        """
        openclaw_status = await self.openclaw_checker.check_all()
        return openclaw_status.to_dict()

    async def _openclaw_check_loop(self):
        """OpenClaw check loop - independent coroutine that sends notifications when issues are detected"""
        logger.info(
            f"[CLIENT-OPENCLAW-CHECK-START] OpenClaw check task started: "
            f"interval={self.openclaw_check_interval}s, client_id={self.client_id}"
        )

        while self.running and self.ws and not self.ws.closed:
            try:
                # Perform OpenClaw check
                status = await self._perform_openclaw_check()
                
                # If issues detected and status changed, send notification
                if status.get("status") != "healthy":
                    # Check if status changed (avoid duplicate notifications)
                    status_key = (status.get("status"), status.get("message"))
                    if status_key != self.last_openclaw_status:
                        await self._send_openclaw_notification(status)
                        self.last_openclaw_status = status_key
                else:
                    # When recovered, if there was a previous issue, also send a notification
                    if self.last_openclaw_status is not None:
                        await self._send_openclaw_notification(status)
                        self.last_openclaw_status = None

                # Wait for next check
                await asyncio.sleep(self.openclaw_check_interval)

            except Exception as e:
                logger.error(
                    f"[CLIENT-OPENCLAW-CHECK-ERROR] OpenClaw check error: {e}, "
                    f"client_id={self.client_id}",
                    exc_info=True
                )
                # Continue running even if check fails
                await asyncio.sleep(self.openclaw_check_interval)

        logger.info(
            f"[CLIENT-OPENCLAW-CHECK-STOP] OpenClaw check task stopped: "
            f"client_id={self.client_id}"
        )

    async def _send_openclaw_notification(self, status: dict):
        """
        Send OpenClaw status notification message
        
        Args:
            status: OpenClaw status dictionary with structure:
                {
                    "status": "healthy|warning|error",
                    "message": "Status description",
                    "details": {
                        # Any check-specific parameters
                    }
                }
        """
        if not self.ws or self.ws.closed:
            logger.warning(
                f"[CLIENT-OPENCLAW-NOTIFY] WebSocket not connected, skipping notification: "
                f"client_id={self.client_id}"
            )
            return

        try:
            notification_msg = {
                "id": str(uuid.uuid4()),
                "type": "health_status",  # Keep "health_status" for protocol compatibility
                "timestamp": int(time.time() * 1000),
                "data": {
                    "status": status.get("status"),
                    "message": status.get("message"),
                    "details": status.get("details", {}),
                }
            }
            
            logger.info(
                f"[CLIENT-OPENCLAW-NOTIFY] Sending OpenClaw notification: "
                f"status={status.get('status')}, message_id={notification_msg['id']}, "
                f"client_id={self.client_id}"
            )
            
            # Send message directly
            await self.ws.send_json(notification_msg)
            
            logger.debug(
                f"[CLIENT-OPENCLAW-NOTIFY] OpenClaw notification sent: "
                f"message_id={notification_msg['id']}, client_id={self.client_id}"
            )
            
        except Exception as e:
            logger.error(
                f"[CLIENT-OPENCLAW-NOTIFY-ERROR] Failed to send OpenClaw notification: "
                f"error={e}, client_id={self.client_id}",
                exc_info=True
            )

    async def _execute_local_command(
        self,
        message_id: str,
        command: str,
        session_id: str,
        args: list = None,
        cwd: Optional[str] = None,
        timeout: int = 300,
    ) -> dict:
        """
        Execute local command

        Args:
            message_id: 消息ID
            command: 命令名称
            session_id: 聊天Session ID（required）
            args: 命令Arguments列表
            cwd: 工作目录
            timeout: 超时时间（seconds）

        Returns:
            dict: 响应消息
        """
        start_time = time.time()
        try:
            # 验证 session_id
            if not session_id:
                logger.error(f"[CLIENT-LOCAL-EXEC-ERROR] session_id is required: message_id={message_id}")
                return {
                    "id": message_id,
                    "type": "local_tool_result",
                    "success": False,
                    "data": {
                        "exit_code": -1,
                        "stdout": "",
                        "stderr": "session_id is required for command execution",
                        "execution_time": 0.0,
                    },
                }

            if args is None:
                args = []

            # 展开Arguments中的 ~ 路径和环境变量
            expanded_args = []
            for arg in args:
                # 只对看起来像路径的Arguments进行展开（包含 ~ 或 $）
                if "~" in arg or "$" in arg:
                    expanded_args.append(os.path.expanduser(os.path.expandvars(arg)))
                else:
                    expanded_args.append(arg)

            # 构建完整命令
            if expanded_args:
                full_command = f"{command} {' '.join(expanded_args)}"
            else:
                full_command = command

            # Whitelist check
            is_allowed, blocked_commands = is_command_allowed(full_command)
            if not is_allowed:
                whitelist_path = get_whitelist_file_path()
                blocked_list = ", ".join(blocked_commands) if blocked_commands else command
                error_msg = (
                    f"Command execution rejected。Following commands not in whitelist: {blocked_list}\n"
                    f"Please add required commands to whitelist configuration file: {whitelist_path}"
                )
                logger.warning(
                    f"[CLIENT-LOCAL-EXEC-BLOCKED] Commands not in whitelist: message_id={message_id}, "
                    f"blocked_commands={blocked_commands}, full_command={full_command[:200]}, "
                    f"whitelist_path={whitelist_path}, client_id={self.client_id}"
                )
                return {
                    "id": message_id,
                    "type": "local_tool_result",
                    "success": False,
                    "data": {
                        "exit_code": -3,
                        "stdout": "",
                        "stderr": error_msg,
                        "execution_time": 0.0,
                    },
                }

            logger.info(
                f"[CLIENT-LOCAL-EXEC-START] Executing local command: message_id={message_id}, "
                f"command={full_command}, session_id={session_id}, cwd={cwd}, timeout={timeout}s, "
                f"client_id={self.client_id}"
            )

            # Create session-specific executor
            executor = CommandExecutor(session_id=session_id, timeout=timeout)

            # Execute command（Pass cwd to execute method if specified）
            result = await executor.execute(full_command, cwd=cwd)

            elapsed = time.time() - start_time
            logger.info(
                f"[CLIENT-LOCAL-EXEC-SUCCESS] Local command completed: message_id={message_id}, "
                f"exit_code={result.exit_code}, elapsed={elapsed:.2f}s, "
                f"stdout_len={len(result.stdout)}, stderr_len={len(result.stderr)}, client_id={self.client_id}"
            )

            return {
                "id": message_id,
                "type": "local_tool_result",
                "success": result.exit_code == 0,
                "data": result.to_dict(),
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(
                f"[CLIENT-LOCAL-EXEC-ERROR] Local command execution failed: message_id={message_id}, "
                f"error={e}, elapsed={elapsed:.2f}s, client_id={self.client_id}",
                exc_info=True,
            )
            return {
                "id": message_id,
                "type": "local_tool_result",
                "success": False,
                "data": {
                    "exit_code": -1,
                    "stdout": "",
                    "stderr": f"Execute local command失败: {str(e)}",
                    "execution_time": 0.0,
                },
            }

    async def _execute_mcp_tool(
        self,
        message_id: str,
        server_name: str,
        tool_name: str,
        session_id: str,
        arguments: dict,
    ) -> dict:
        """
        Execute MCP tool

        Args:
            message_id: 消息ID
            server_name: MCP 服务器名称
            tool_name: 工具名称
            session_id: 聊天Session ID（required）
            arguments: 工具Arguments

        Returns:
            dict: 响应消息
        """
        start_time = time.time()
        try:
            # 验证 session_id
            if not session_id:
                logger.error(f"[CLIENT-MCP-EXEC-ERROR] session_id is required: message_id={message_id}")
                return {
                    "id": message_id,
                    "type": "mcp_tool_result",
                    "success": False,
                    "data": {
                        "server_name": server_name,
                        "tool_name": tool_name,
                        "result": None,
                        "error": "session_id is required for MCP tool execution",
                    },
                }

            logger.info(
                f"[CLIENT-MCP-EXEC-START] Executing MCP tool: message_id={message_id}, "
                f"server={server_name}, tool={tool_name}, session_id={session_id}, arguments={arguments}, "
                f"client_id={self.client_id}"
            )

            # Call MCP tool
            if not self.mcp_manager.client:
                logger.error(
                    f"[CLIENT-MCP-EXEC-ERROR] MCP client not initialized: message_id={message_id}, "
                    f"client_id={self.client_id}"
                )
                return {
                    "id": message_id,
                    "type": "mcp_tool_result",
                    "success": False,
                    "data": {
                        "server_name": server_name,
                        "tool_name": tool_name,
                        "result": None,
                        "error": "MCP 客户端未初始化",
                    },
                }

            # 执行工具
            tools = await self.mcp_manager.client.get_tools(server_name=server_name)
            current_tool = None
            # Iterate through tool_name to get corresponding tool
            for tool in tools:
                if tool.name == tool_name:
                    current_tool = tool
                    break
            if not current_tool:
                logger.error(
                    f"[CLIENT-MCP-EXEC-ERROR] MCP tool not found: message_id={message_id}, "
                    f"server={server_name}, tool={tool_name}, available_tools={[t.name for t in tools]}, "
                    f"client_id={self.client_id}"
                )
                return {
                    "id": message_id,
                    "type": "mcp_tool_result",
                    "success": False,
                    "data": {
                        "server_name": server_name,
                        "tool_name": tool_name,
                        "result": None,
                        "error": "Execute MCP tool失败: tool不存在",
                    },
                }

            result = await current_tool.ainvoke(input=arguments)

            elapsed = time.time() - start_time
            logger.info(
                f"[CLIENT-MCP-EXEC-SUCCESS] MCP tool completed: message_id={message_id}, "
                f"server={server_name}, tool={tool_name}, elapsed={elapsed:.2f}s, client_id={self.client_id}"
            )

            # Log to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="MCP_TOOL",
                operation=f"{server_name}:{tool_name}",
                arguments=arguments,
                result=result,
                success=True,
                elapsed=elapsed,
            )

            return {
                "id": message_id,
                "type": "mcp_tool_result",
                "success": True,
                "data": {
                    "server_name": server_name,
                    "tool_name": tool_name,
                    "result": result,
                },
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(
                f"[CLIENT-MCP-EXEC-ERROR] MCP tool execution failed: message_id={message_id}, "
                f"server={server_name}, tool={tool_name}, error={e}, elapsed={elapsed:.2f}s, "
                f"client_id={self.client_id}",
                exc_info=True,
            )
            return {
                "id": message_id,
                "type": "mcp_tool_result",
                "success": False,
                "data": {
                    "server_name": server_name,
                    "tool_name": tool_name,
                    "result": None,
                    "error": f"Execute MCP tool失败: {str(e)}",
                },
            }

    async def _execute_doc_read(
        self,
        message_id: str,
        file_path: str,
        session_id: str,
        encoding: Optional[str] = None,
    ) -> dict:
        """
        Execute document read

        Args:
            message_id: 消息ID
            file_path: 文件路径
            session_id: 聊天Session ID
            encoding: 文件编码

        Returns:
            dict: 响应消息
        """
        start_time = time.time()
        try:
            logger.info(
                f"[CLIENT-DOC-READ-EXEC-START] Executing doc read: message_id={message_id}, "
                f"file_path={file_path}, session_id={session_id}, encoding={encoding}, client_id={self.client_id}"
            )

            # Call document.read_document
            result = document.read_document(file_path=file_path, encoding=encoding)

            elapsed = time.time() - start_time
            logger.info(
                f"[CLIENT-DOC-READ-EXEC-SUCCESS] Doc read completed: message_id={message_id}, "
                f"success={result.get('success')}, elapsed={elapsed:.2f}s, client_id={self.client_id}"
            )

            # Log to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_READ",
                operation=file_path,
                arguments={"encoding": encoding},
                result={"size": result.get("size"), "encoding": result.get("encoding")},
                success=result.get("success", False),
                elapsed=elapsed,
                error=result.get("error"),
            )

            return {
                "id": message_id,
                "type": "doc_read_result",
                "success": result.get("success", False),
                "data": {
                    "content": result.get("content"),
                    "encoding": result.get("encoding"),
                    "size": result.get("size"),
                    "error": result.get("error"),
                },
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(
                f"[CLIENT-DOC-READ-EXEC-ERROR] Doc read execution failed: message_id={message_id}, "
                f"error={e}, elapsed={elapsed:.2f}s, client_id={self.client_id}",
                exc_info=True,
            )

            # Log error to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_READ",
                operation=file_path,
                arguments={"encoding": encoding},
                result=None,
                success=False,
                elapsed=elapsed,
                error=str(e),
            )

            return {
                "id": message_id,
                "type": "doc_read_result",
                "success": False,
                "data": {
                    "content": None,
                    "encoding": None,
                    "size": None,
                    "error": f"Execute document read失败: {str(e)}",
                },
            }

    async def _execute_doc_write(
        self,
        message_id: str,
        file_path: str,
        session_id: str,
        content: str,
        encoding: str = "utf-8",
        create_dirs: bool = True,
    ) -> dict:
        """
        Execute document write

        Args:
            message_id: 消息ID
            file_path: 文件路径
            session_id: 聊天Session ID
            content: 文件内容
            encoding: 文件编码
            create_dirs: 是否创建父目录

        Returns:
            dict: 响应消息
        """
        start_time = time.time()
        try:
            logger.info(
                f"[CLIENT-DOC-WRITE-EXEC-START] Executing doc write: message_id={message_id}, "
                f"file_path={file_path}, content_len={len(content)}, session_id={session_id}, encoding={encoding}, "
                f"create_dirs={create_dirs}, client_id={self.client_id}"
            )

            # Call document.write_document
            result = document.write_document(
                file_path=file_path,
                content=content,
                encoding=encoding,
                create_dirs=create_dirs,
            )

            elapsed = time.time() - start_time
            logger.info(
                f"[CLIENT-DOC-WRITE-EXEC-SUCCESS] Doc write completed: message_id={message_id}, "
                f"success={result.get('success')}, elapsed={elapsed:.2f}s, client_id={self.client_id}"
            )

            # Log to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_WRITE",
                operation=file_path,
                arguments={
                    "content_length": len(content),
                    "encoding": encoding,
                    "create_dirs": create_dirs,
                },
                result={"size": result.get("size"), "path": result.get("path")},
                success=result.get("success", False),
                elapsed=elapsed,
                error=result.get("error"),
            )

            return {
                "id": message_id,
                "type": "doc_write_result",
                "success": result.get("success", False),
                "data": {
                    "path": result.get("path"),
                    "size": result.get("size"),
                    "error": result.get("error"),
                },
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(
                f"[CLIENT-DOC-WRITE-EXEC-ERROR] Doc write execution failed: message_id={message_id}, "
                f"error={e}, elapsed={elapsed:.2f}s, client_id={self.client_id}",
                exc_info=True,
            )

            # Log error to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_WRITE",
                operation=file_path,
                arguments={
                    "content_length": len(content),
                    "encoding": encoding,
                    "create_dirs": create_dirs,
                },
                result=None,
                success=False,
                elapsed=elapsed,
                error=str(e),
            )

            return {
                "id": message_id,
                "type": "doc_write_result",
                "success": False,
                "data": {
                    "path": None,
                    "size": None,
                    "error": f"Execute document write失败: {str(e)}",
                },
            }

    async def _execute_doc_edit(
        self,
        message_id: str,
        file_path: str,
        session_id: str,
        operation: str,
        new_content: str,
        old_content: Optional[str] = None,
        encoding: Optional[str] = None,
    ) -> dict:
        """
        Execute document edit

        Args:
            message_id: 消息ID
            file_path: 文件路径
            session_id: 聊天Session ID
            operation: Operation type（replace, append, prepend）
            new_content: 新内容
            old_content: 旧内容（Only needed for replace operation）
            encoding: 文件编码

        Returns:
            dict: 响应消息
        """
        start_time = time.time()
        try:
            logger.info(
                f"[CLIENT-DOC-EDIT-EXEC-START] Executing doc edit: message_id={message_id}, "
                f"file_path={file_path}, operation={operation}, session_id={session_id}, encoding={encoding}, client_id={self.client_id}"
            )

            # Call document.edit_document
            result = document.edit_document(
                file_path=file_path,
                operation=operation,
                new_content=new_content,
                old_content=old_content,
                encoding=encoding,
            )

            elapsed = time.time() - start_time
            logger.info(
                f"[CLIENT-DOC-EDIT-EXEC-SUCCESS] Doc edit completed: message_id={message_id}, "
                f"success={result.get('success')}, operation={operation}, elapsed={elapsed:.2f}s, "
                f"client_id={self.client_id}"
            )

            # Log to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_EDIT",
                operation=f"{operation}:{file_path}",
                arguments={
                    "new_content_length": len(new_content),
                    "old_content_length": len(old_content) if old_content else 0,
                    "encoding": encoding,
                },
                result={"size": result.get("size"), "path": result.get("path")},
                success=result.get("success", False),
                elapsed=elapsed,
                error=result.get("error"),
            )

            return {
                "id": message_id,
                "type": "doc_edit_result",
                "success": result.get("success", False),
                "data": {
                    "operation": result.get("operation"),
                    "size": result.get("size"),
                    "path": result.get("path"),
                    "error": result.get("error"),
                },
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(
                f"[CLIENT-DOC-EDIT-EXEC-ERROR] Doc edit execution failed: message_id={message_id}, "
                f"error={e}, elapsed={elapsed:.2f}s, client_id={self.client_id}",
                exc_info=True,
            )

            # Log error to terminal.log
            self._log_operation(
                session_id=session_id,
                operation_type="DOC_EDIT",
                operation=f"{operation}:{file_path}",
                arguments={
                    "new_content_length": len(new_content),
                    "old_content_length": len(old_content) if old_content else 0,
                    "encoding": encoding,
                },
                result=None,
                success=False,
                elapsed=elapsed,
                error=str(e),
            )

            return {
                "id": message_id,
                "type": "doc_edit_result",
                "success": False,
                "data": {
                    "operation": operation,
                    "size": None,
                    "path": None,
                    "error": f"Execute document edit失败: {str(e)}",
                },
            }

    async def _listen(self):
        """监听Server message"""
        try:
            logger.info(f"[CLIENT-LISTEN-START] Started listening for messages: client_id={self.client_id}")

            async for msg in self.ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        # Parse message
                        data = json.loads(msg.data)
                        logger.debug(f"[CLIENT-WS-RECV] Received raw message: data={data}, client_id={self.client_id}")

                        # Put command processing in background task to avoid blocking message receive loop
                        # This ensures other messages (like ping/pong) can still be received and processed during long tasks
                        asyncio.create_task(self._handle_and_respond(data))

                    except json.JSONDecodeError as e:
                        logger.error(
                            f"[CLIENT-MSG-PARSE-ERROR] Failed to parse message: error={e}, "
                            f"raw_data={msg.data}, client_id={self.client_id}"
                        )
                    except Exception as e:
                        logger.error(
                            f"[CLIENT-MSG-HANDLE-ERROR] Error handling message: error={e}, client_id={self.client_id}",
                            exc_info=True,
                        )

                elif msg.type == aiohttp.WSMsgType.ERROR:
                    logger.error(
                        f"[CLIENT-WS-ERROR] WebSocket error: exception={self.ws.exception()}, "
                        f"client_id={self.client_id}"
                    )
                    break

                elif msg.type == aiohttp.WSMsgType.CLOSED:
                    logger.info(f"[CLIENT-WS-CLOSED] WebSocket connection closed: client_id={self.client_id}")
                    break

        except Exception as e:
            logger.error(
                f"[CLIENT-LISTEN-ERROR] Error in message listener: error={e}, client_id={self.client_id}",
                exc_info=True,
            )

    async def _handle_and_respond(self, data: dict):
        """
        Handle command and send response (run as background task)

        Put command handling and response sending in separate task to avoid blocking _listen() loop。
        这样可以确保在长任务执行期间：
        1. Message receive loop continues running, can receive server PING/PONG messages
        2. Heartbeat task can send normally
        3. WebSocket Connection stays active

        Args:
            data: Received message data
        """
        try:
            # Handle command
            response = await self._handle_command(data)

            # Send response（if any）
            if response is not None:
                logger.info(
                    f"[CLIENT-RESPONSE-SEND] Sending response: message_id={response.get('id')}, "
                    f"type={response.get('type')}, success={response.get('success')}, "
                    f"client_id={self.client_id}"
                )
                await self.ws.send_json(response)
                logger.info(
                    f"[CLIENT-RESPONSE-SENT] Response sent successfully: message_id={response.get('id')}, "
                    f"client_id={self.client_id}"
                )
            else:
                logger.debug(
                    f"[CLIENT-NO-RESPONSE] No response needed for message: "
                    f"message_id={data.get('id')}, type={data.get('type')}, client_id={self.client_id}"
                )
        except Exception as e:
            logger.error(
                f"[CLIENT-HANDLE-RESPOND-ERROR] Error in handle_and_respond: error={e}, "
                f"message_id={data.get('id')}, client_id={self.client_id}",
                exc_info=True,
            )

    async def connect(self):
        """Establish WebSocket connection"""
        ws_url = self._get_ws_url()
        logger.info(
            f"[CLIENT-CONNECT-START] Attempting to connect to server: url={ws_url}, "
            f"client_id={self.client_id}, workspace_id={self.workspace_id}"
        )

        try:
            # Don't set global timeout when creating ClientSession
            # This ensures long-running commands won't be interrupted by timeout
            self.session = aiohttp.ClientSession()
            logger.debug(f"[CLIENT-SESSION-CREATED] aiohttp session created: client_id={self.client_id}")

            # Establish WebSocket connection，配置适合长时间任务的Arguments
            self.ws = await self.session.ws_connect(
                ws_url,
                heartbeat=self.heartbeat_interval,  # Enable WebSocket protocol-level heartbeat（PING/PONG）
                timeout=900.0,  # 15minute timeout，Support long-running commands
                autoclose=False,  # Don't auto-close connection
                autoping=True,  # Auto-respond to server PING
            )
            logger.info(
                f"[CLIENT-CONNECT-SUCCESS] WebSocket connection established: client_id={self.client_id}, "
                f"ws_closed={self.ws.closed}, heartbeat={self.heartbeat_interval}s"
            )

            # Start heartbeat task
            self.heartbeat_task = asyncio.create_task(self._send_heartbeat())
            logger.info(
                f"[CLIENT-HEARTBEAT-TASK] Heartbeat task started: interval={self.heartbeat_interval}s, "
                f"client_id={self.client_id}"
            )

            # Start OpenClaw check task (if enabled)
            if self.openclaw_check_enabled:
                self.openclaw_check_task = asyncio.create_task(self._openclaw_check_loop())
                logger.info(
                    f"[CLIENT-OPENCLAW-CHECK-TASK] OpenClaw check task started: "
                    f"interval={self.openclaw_check_interval}s, client_id={self.client_id}"
                )
            else:
                logger.debug(
                    f"[CLIENT-OPENCLAW-CHECK-TASK] OpenClaw check disabled, skipping task start: "
                    f"client_id={self.client_id}"
                )

            return True

        except Exception as e:
            logger.error(
                f"[CLIENT-CONNECT-ERROR] Failed to connect: error={e}, client_id={self.client_id}",
                exc_info=True,
            )
            if self.session:
                await self.session.close()
            return False

    async def disconnect(self):
        """Disconnect WebSocket connection"""
        logger.info(f"[CLIENT-DISCONNECT-START] Starting disconnect process: client_id={self.client_id}")

        # Stop OpenClaw check task
        if self.openclaw_check_task and not self.openclaw_check_task.done():
            logger.debug(
                f"[CLIENT-DISCONNECT-OPENCLAW-CHECK] Cancelling OpenClaw check task: "
                f"client_id={self.client_id}"
            )
            self.openclaw_check_task.cancel()
            try:
                await self.openclaw_check_task
            except asyncio.CancelledError:
                pass
            logger.info(
                f"[CLIENT-DISCONNECT-OPENCLAW-CHECK] OpenClaw check task stopped: "
                f"client_id={self.client_id}"
            )

        # Stop heartbeat task
        if self.heartbeat_task and not self.heartbeat_task.done():
            logger.debug(f"[CLIENT-DISCONNECT-HEARTBEAT] Cancelling heartbeat task: client_id={self.client_id}")
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass
            logger.info(f"[CLIENT-DISCONNECT-HEARTBEAT] Heartbeat task stopped: client_id={self.client_id}")

        if self.ws and not self.ws.closed:
            logger.debug(f"[CLIENT-DISCONNECT-WS] Closing WebSocket connection: client_id={self.client_id}")
            await self.ws.close()
            logger.info(f"[CLIENT-DISCONNECT-WS] WebSocket closed: client_id={self.client_id}")

        if self.session:
            logger.debug(f"[CLIENT-DISCONNECT-SESSION] Closing aiohttp session: client_id={self.client_id}")
            await self.session.close()
            logger.info(f"[CLIENT-DISCONNECT-SESSION] Session closed: client_id={self.client_id}")

        logger.info(f"[CLIENT-DISCONNECT-COMPLETE] Disconnect completed: client_id={self.client_id}")

    async def run(self):
        """Run client (with automatic reconnection)"""
        self.running = True
        logger.info(
            f"[CLIENT-RUN-START] Starting client: client_id={self.client_id}, "
            f"server_url={self.server_url}, workspace_id={self.workspace_id}"
        )

        # Connect to MCP services
        try:
            logger.info(f"[CLIENT-MCP-CONNECT] Connecting to MCP services: client_id={self.client_id}")
            mcp_count = await self.mcp_manager.connect_all()
            if mcp_count > 0:
                logger.info(
                    f"[CLIENT-MCP-CONNECTED] Successfully connected to MCP services: count={mcp_count}, "
                    f"client_id={self.client_id}"
                )
            else:
                logger.warning(f"[CLIENT-MCP-NONE] No MCP services configured or connected: client_id={self.client_id}")
        except Exception as e:
            logger.error(
                f"[CLIENT-MCP-ERROR] Error connecting to MCP services: error={e}, client_id={self.client_id}",
                exc_info=True,
            )

        reconnect_count = 0
        while self.running:
            # Attempting to connect
            logger.info(
                f"[CLIENT-RECONNECT] Attempting to connect: attempt={reconnect_count + 1}, client_id={self.client_id}"
            )

            if await self.connect():
                reconnect_count = 0  # Reset reconnection count
                try:
                    # Listen for messages
                    await self._listen()
                except Exception as e:
                    logger.error(
                        f"[CLIENT-RUN-ERROR] Error during client run: error={e}, client_id={self.client_id}",
                        exc_info=True,
                    )
                finally:
                    # Disconnect
                    await self.disconnect()
            else:
                reconnect_count += 1
                logger.warning(
                    f"[CLIENT-CONNECT-FAILED] Failed to establish connection: attempts={reconnect_count}, "
                    f"client_id={self.client_id}"
                )

            # If still running, wait then reconnect
            if self.running:
                logger.info(
                    f"[CLIENT-RECONNECT-WAIT] Waiting for reconnection: delay={self.reconnect_interval}s, "
                    f"client_id={self.client_id}"
                )
                await asyncio.sleep(self.reconnect_interval)

    async def stop(self):
        """Stop client"""
        logger.info(f"[CLIENT-STOP-START] Stopping client: client_id={self.client_id}")

        self.running = False
        await self.disconnect()

        # Disconnect from MCP services
        try:
            logger.info(f"[CLIENT-MCP-DISCONNECT] Disconnecting from MCP services: client_id={self.client_id}")
            await self.mcp_manager.disconnect_all()
            logger.info(f"[CLIENT-MCP-DISCONNECTED] MCP services disconnected: client_id={self.client_id}")
        except Exception as e:
            logger.error(
                f"[CLIENT-MCP-DISCONNECT-ERROR] Error disconnecting from MCP services: error={e}, "
                f"client_id={self.client_id}",
                exc_info=True,
            )

        logger.info(f"[CLIENT-STOP-COMPLETE] Client stopped: client_id={self.client_id}")
