#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：deploy_executor.py
模块描述：自动化部署 SSH 执行器，通过 paramiko 连接远程服务器执行部署脚本

作者：czj
创建日期：2026-08-12
最后修改：2026-08-12
"""

import asyncio
import logging
import re
import time
from io import StringIO
from typing import Optional, Callable, List, Dict, Any

import paramiko

logger = logging.getLogger(__name__)


class DeployExecutor:
    """
    自动化部署 SSH 执行器

    通过 paramiko 建立 SSH 连接，在远程服务器上执行部署脚本，
    支持交互式应答（预设触发词→应答映射）和实时日志推送。

    Attributes:
        host (str): 服务器地址
        port (int): SSH 端口
        username (str): 登录用户名
        password (Optional[str]): 密码（密钥认证时为 None）
        private_key (Optional[str]): 私钥内容（密码认证时为 None）
        timeout (int): 执行超时时间（秒）
        expect_answers (List[Dict]): 交互式应答规则列表
        _client (Optional[paramiko.SSHClient]): SSH 客户端
        _channel (Optional[paramiko.Channel]): SSH 通道
        _cancelled (bool): 是否已取消
        _ev_loop (Optional[asyncio.AbstractEventLoop]): 主事件循环引用
        _broadcast_fn (Optional[Callable]): WebSocket 广播函数引用
    """

    def __init__(
        self,
        host: str,
        port: int = 22,
        username: str = "",
        password: Optional[str] = None,
        private_key: Optional[str] = None,
        timeout: int = 3600,
        expect_answers: Optional[List[Dict[str, Any]]] = None,
    ):
        """
        初始化 SSH 执行器

        Args:
            host: 服务器地址
            port: SSH 端口，默认 22
            username: 登录用户名
            password: SSH 密码（密钥认证时传 None）
            private_key: SSH 私钥内容（密码认证时传 None）
            timeout: 执行超时时间（秒）
            expect_answers: 交互式应答规则列表
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.private_key = private_key
        self.timeout = timeout
        self.expect_answers = expect_answers or []

        self._client: Optional[paramiko.SSHClient] = None
        self._channel: Optional[paramiko.Channel] = None
        self._cancelled = False
        self._ev_loop: Optional[asyncio.AbstractEventLoop] = None
        self._broadcast_fn: Optional[Callable] = None

    # ==================== 广播回调设置 ====================

    def set_broadcast(
        self,
        ev_loop: asyncio.AbstractEventLoop,
        broadcast_fn: Callable,
    ) -> None:
        """
        设置广播回调（由服务层调用，用于跨线程推送日志）

        Args:
            ev_loop: 主事件循环
            broadcast_fn: 异步广播函数（如 manager.broadcast）
        """
        self._ev_loop = ev_loop
        self._broadcast_fn = broadcast_fn

    def _broadcast(self, execution_id: str, message: Dict[str, Any]) -> None:
        """
        线程安全地广播消息到 WebSocket

        通过 asyncio.run_coroutine_threadsafe 将异步调用提交到主事件循环。

        Args:
            execution_id: 执行记录 ID
            message: 消息内容
        """
        if self._ev_loop and self._broadcast_fn:
            try:
                asyncio.run_coroutine_threadsafe(
                    self._broadcast_fn(execution_id, message),
                    self._ev_loop,
                )
            except Exception as e:
                logger.error(f"[DeployExecutor] 广播失败: {e}")

    # ==================== SSH 连接管理 ====================

    def _connect(self) -> None:
        """
        建立 SSH 连接

        支持密码认证和密钥认证两种方式。
        私钥自动尝试 RSA 和 Ed25519 两种格式。

        Raises:
            paramiko.AuthenticationException: 认证失败
            paramiko.SSHException: SSH 协议错误
            Exception: 其他连接错误
        """
        self._client = paramiko.SSHClient()
        self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        connect_kwargs: Dict[str, Any] = {
            "hostname": self.host,
            "port": self.port,
            "username": self.username,
            "timeout": 30,
            "banner_timeout": 30,
        }

        if self.private_key:
            # 密钥认证
            key_file = StringIO(self.private_key)
            try:
                private_key_obj = paramiko.RSAKey.from_private_key(key_file)
            except paramiko.SSHException:
                key_file.seek(0)
                try:
                    private_key_obj = paramiko.Ed25519Key.from_private_key(key_file)
                except paramiko.SSHException:
                    key_file.seek(0)
                    private_key_obj = paramiko.ECDSAKey.from_private_key(key_file)
            connect_kwargs["pkey"] = private_key_obj
        else:
            # 密码认证
            connect_kwargs["password"] = self.password

        self._client.connect(**connect_kwargs)
        logger.info(
            f"[DeployExecutor] SSH 连接成功: {self.username}@{self.host}:{self.port}"
        )

    def _open_shell(self) -> None:
        """
        打开交互式 PTY 会话

        使用 invoke_shell + get_pty 开启 PTY 模式。
        脚本依赖 isatty() 检测时必需使用 PTY，否则部分脚本可能行为异常。
        """
        self._channel = self._client.invoke_shell(term="xterm")
        self._channel.get_pty()
        self._channel.settimeout(0.5)  # 非阻塞读取

        # 等待 shell 就绪，丢弃初始 banner
        if self.expect_answers:
            time.sleep(0.5)
            try:
                self._channel.recv(4096)
            except Exception:
                pass

    # ==================== 核心执行逻辑 ====================

    def execute(
        self,
        execution_id: str,
        command: str,
        work_dir: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        执行部署脚本并处理交互式应答

        Args:
            execution_id: 执行记录 ID
            command: 远程执行命令
            work_dir: 远程工作目录（可选，会在命令前先 cd）

        Returns:
            Dict: 执行结果，包含 exit_code、error_message、duration

        Raises:
            RuntimeError: 执行器未连接
        """
        self._cancelled = False
        start_time = time.time()

        # ---- 建立连接 ----
        try:
            self._connect()
            self._open_shell()
        except paramiko.AuthenticationException:
            return {
                "exit_code": -1,
                "error_message": "SSH 认证失败，请检查用户名和密码/密钥",
                "duration": round(time.time() - start_time, 2),
            }
        except paramiko.SSHException as e:
            return {
                "exit_code": -1,
                "error_message": f"SSH 连接异常: {str(e)}",
                "duration": round(time.time() - start_time, 2),
            }
        except OSError as e:
            return {
                "exit_code": -1,
                "error_message": f"网络连接失败: {str(e)}",
                "duration": round(time.time() - start_time, 2),
            }
        except Exception as e:
            logger.error(f"[DeployExecutor] 连接失败: {e}")
            return {
                "exit_code": -1,
                "error_message": f"连接失败: {str(e)}",
                "duration": round(time.time() - start_time, 2),
            }

        # ---- 广播连接成功 ----
        self._broadcast(execution_id, {
            "type": "deploy_start",
            "data": {
                "execution_id": execution_id,
                "message": f"已连接到 {self.username}@{self.host}:{self.port}",
                "timestamp": time.time(),
            },
        })

        # ---- 构造并发送命令 ----
        if work_dir:
            full_command = f"cd {work_dir} && {command}"
        else:
            full_command = command

        self._channel.send(full_command + "\n")
        self._broadcast(execution_id, {
            "type": "deploy_output",
            "data": {
                "execution_id": execution_id,
                "message": f"$ {full_command}\n",
                "timestamp": time.time(),
            },
        })

        # ---- 单线程轮询循环 ----
        buffer = ""
        matched_indices: set = set()  # 已匹配的规则索引
        exit_code = None
        error_message = None
        deadline = time.time() + self.timeout

        # 按 order 排序的应答规则（带原始索引）
        sorted_answers = sorted(
            enumerate(self.expect_answers),
            key=lambda x: x[1].get("order", 0),
        )

        while not self._cancelled:
            # 超时检查
            if time.time() > deadline:
                error_message = f"执行超时（{self.timeout}秒）"
                logger.warning(f"[DeployExecutor] {error_message}")
                break

            # ---- 读取输出 ----
            try:
                if self._channel and self._channel.recv_ready():
                    data = self._channel.recv(4096).decode("utf-8", errors="replace")
                    buffer += data

                    # 敏感信息过滤（日志中隐藏密码等）
                    display_msg = data
                    for _idx, rule in sorted_answers:
                        if rule.get("sensitive") and rule.get("answer"):
                            display_msg = display_msg.replace(
                                rule["answer"], "******"
                            )

                    self._broadcast(execution_id, {
                        "type": "deploy_output",
                        "data": {
                            "execution_id": execution_id,
                            "message": display_msg,
                            "timestamp": time.time(),
                        },
                    })
            except paramiko.SSHException:
                # 通道可能已关闭，退出读取循环
                break
            except Exception as e:
                logger.error(f"[DeployExecutor] 读取输出异常: {e}")
                break

            # ---- 检查退出状态 ----
            if self._channel and self._channel.exit_status_ready():
                exit_code = self._channel.recv_exit_status()
                logger.info(
                    f"[DeployExecutor] 命令完成，退出码: {exit_code}"
                )
                break

            # ---- 匹配交互式应答规则 ----
            for idx, rule in sorted_answers:
                if idx in matched_indices:
                    continue

                trigger = rule.get("trigger", "")
                trigger_type = rule.get("trigger_type", "contains")
                answer = rule.get("answer", "")

                matched = False
                if trigger_type == "regex":
                    try:
                        if re.search(trigger, buffer, re.MULTILINE):
                            matched = True
                    except re.error:
                        logger.warning(
                            f"[DeployExecutor] 正则表达式错误: {trigger}"
                        )
                else:  # contains
                    if trigger in buffer:
                        matched = True

                if matched:
                    # 发送应答
                    self._channel.send(answer + "\n")
                    matched_indices.add(idx)

                    # 广播应答（敏感信息隐藏）
                    display_answer = (
                        "******" if rule.get("sensitive") else answer
                    )
                    self._broadcast(execution_id, {
                        "type": "deploy_answer_sent",
                        "data": {
                            "execution_id": execution_id,
                            "trigger": trigger,
                            "answer": display_answer,
                            "sensitive": rule.get("sensitive", False),
                            "timestamp": time.time(),
                        },
                    })
                    logger.info(
                        f"[DeployExecutor] 应答已发送: "
                        f"trigger='{trigger}' answer='{display_answer}'"
                    )
                    break  # 每轮只匹配一条规则，避免竞态

            # 无数据时短暂休眠，避免 CPU 忙等
            if not (self._channel and self._channel.recv_ready()):
                time.sleep(0.05)

        # ---- 执行结束 ----
        duration = round(time.time() - start_time, 2)

        if self._cancelled:
            exit_code = -1
            error_message = "执行已被用户取消"

        self.close()

        return {
            "exit_code": exit_code,
            "error_message": error_message,
            "duration": duration,
        }

    # ==================== 取消与清理 ====================

    def cancel(self) -> None:
        """
        取消当前执行

        设置取消标记并关闭通道，强制退出执行循环。
        调用后 execute() 会尽快返回。
        """
        logger.info("[DeployExecutor] 正在取消执行...")
        self._cancelled = True
        if self._channel:
            try:
                self._channel.close()
            except Exception:
                pass

    def close(self) -> None:
        """
        关闭 SSH 连接，释放资源

        安全关闭通道和客户端连接，确保资源不泄漏。
        """
        if self._channel:
            try:
                self._channel.close()
            except Exception:
                pass
            self._channel = None

        if self._client:
            try:
                self._client.close()
            except Exception:
                pass
            self._client = None

        logger.info("[DeployExecutor] SSH 连接已关闭")