"""AI Agent 可用的命令工具和安全策略。

支持两类命令：
1. kubectl 命令 - 通过 classify_kubectl_command 分类
2. 系统诊断 shell 命令 - 通过 classify_shell_command 分类 (白名单)

classify_agent_command 统一分发，execute_command 自动选择执行方式。
所有命令均通过 subprocess_exec 执行 (无 shell)，禁止管道、重定向和命令组合。
"""
import asyncio
import shlex
from dataclasses import dataclass

from backend.diagnosis import KUBECTL

# ── kubectl 命令分类 ──────────────────────────────────────

READ_ONLY_VERBS = {
    "api-resources", "api-versions", "auth", "cluster-info", "config", "describe",
    "diff", "explain", "get", "logs", "top", "version", "wait",
}
WRITE_OR_INTERACTIVE_VERBS = {
    "annotate", "apply", "attach", "autoscale", "cordon", "cp", "create",
    "debug", "delete", "drain", "edit", "exec", "expose", "label", "patch",
    "port-forward", "replace", "rollout", "run", "scale", "set", "taint",
    "uncordon",
}
GLOBAL_FLAGS_WITH_VALUE = {
    "--as", "--as-group", "--as-uid", "--cache-dir", "--certificate-authority",
    "--client-certificate", "--client-key", "--cluster", "--context", "--kubeconfig",
    "--namespace", "-n", "--profile", "--request-timeout", "--server", "-s",
    "--tls-server-name", "--token", "--user",
}
FORBIDDEN_TOKENS = {";", "||", "&&", "<", "<<", "`"}
# 管道 |、重定向 >、>>、$ 变量引用 允许在诊断中使用
# 但含 | 的命令中若管道后出现 sh/bash/dd/format/mkfs/fdisk/parted 则直接拒绝（武器级危险）
# rm/mv/chmod/chown — 现已在 SHELL_WRITE_COMMANDS 中，管道后会触发审批而非直接拒绝
_PIPE_DANGEROUS_TARGETS = {"sh", "bash", "dash", "zsh",
                           "dd", "format", "mkfs", "fdisk", "parted"}


@dataclass(frozen=True)
class CommandDecision:
    mode: str          # "read" (只读，自动执行) 或 "write" (修改，需审批)
    requires_approval: bool
    verb: str          # kubectl 子命令 或 shell 命令名
    args: list[str]
    needs_shell: bool = False  # True 时需用 shell=True 执行


def _has_shell_operators(command: str) -> bool:
    """检测命令是否包含管道 | 或重定向 >。"""
    # 简单检测：提取所有 token 中的操作符
    # 注意排除出现在引号内的 | 和 >
    stripped = command.strip()
    if not stripped:
        return False
    # 先用 shlex 分割看是否含有 shell 操作符作为独立 token
    # 但更可靠的是直接扫描字符（考虑引号保护）
    in_single = in_double = False
    for i, ch in enumerate(stripped):
        if ch == "'" and not in_double:
            in_single = not in_single
        elif ch == '"' and not in_single:
            in_double = not in_double
        if in_single or in_double:
            continue
        if ch in ('|', '>', ';', '&'):
            # ; 和 && 仍然不允许
            if ch in (';', '&'):
                # 单独 & 不算（后台），&& 才不允许
                if ch == ';':
                    return True
                # 检查是否是 &&
                if ch == '&' and i + 1 < len(stripped) and stripped[i + 1] == '&':
                    return True
            if ch == '|':
                # | 允许；但检查是否是 ||
                if i + 1 < len(stripped) and stripped[i + 1] == '|':
                    # || 不允许
                    return False
                return True  # 单 | 是管道，允许
            if ch == '>':
                return True  # > 或 >> 允许
    return False


def _pipe_to_dangerous_target(command: str) -> bool:
    """检查管道后面的命令是否是危险命令。"""
    if '|' not in command:
        return False
    # 按管道分割，取最后一段
    # 简单解析：用 | 分割，忽略引号内的
    parts = []
    buf = []
    in_single = in_double = False
    for ch in command:
        if ch == "'" and not in_double:
            in_single = not in_single
        elif ch == '"' and not in_single:
            in_double = not in_double
        if ch == '|' and not in_single and not in_double:
            parts.append(''.join(buf))
            buf = []
        else:
            buf.append(ch)
    parts.append(''.join(buf))
    if len(parts) < 2:
        return False
    # 检查管道最后一条命令的危险词
    last_part = parts[-1].strip()
    # shlex 取第一个词
    try:
        last_args = shlex.split(last_part)
    except ValueError:
        return False
    if last_args and last_args[0] in _PIPE_DANGEROUS_TARGETS:
        return True
    # 也检查中间管道（如 docker exec | sh）
    for part in parts[1:]:
        try:
            part_args = shlex.split(part.strip())
        except ValueError:
            continue
        if part_args and part_args[0] in _PIPE_DANGEROUS_TARGETS:
            return True
    return False


def _parse_command(command: str) -> list[str]:
    """解析命令，检查合法性。

    允许 kubectl 和诊断命令中的管道 (|)、重定向 (>)、$ 变量引用。
    不允许的: ; || && < << 后台命令 ` 以及多行。
    管道后不允许出现 sh/bash/rm/mv/chown/chmod/dd 等危险命令。
    """
    if not isinstance(command, str) or not command.strip():
        raise ValueError("命令不能为空")
    if "\n" in command or "\r" in command:
        raise ValueError("不允许多行命令")
    if any(token in command for token in FORBIDDEN_TOKENS):
        raise ValueError("不允许 shell 命令连接符（; || &&）或输入重定向（< <<）")
    if _pipe_to_dangerous_target(command):
        raise ValueError("管道后不允许执行危险命令（sh/rm/chmod 等），如需要请手动执行")
    try:
        return shlex.split(command)
    except ValueError as exc:
        raise ValueError(f"命令格式错误: {exc}") from exc


def parse_kubectl_command(command: str) -> list[str]:
    """仅解析单条 kubectl 命令，不允许 shell、重定向或命令组合。"""
    parts = _parse_command(command)
    if not parts or parts[0] != "kubectl":
        raise ValueError("只允许执行 kubectl 命令")
    if len(parts) < 2:
        raise ValueError("kubectl 命令缺少操作类型")
    return parts[1:]


def _find_verb(args: list[str]) -> str:
    index = 0
    while index < len(args):
        arg = args[index]
        if arg == "--":
            break
        if arg in GLOBAL_FLAGS_WITH_VALUE:
            index += 2
            continue
        if any(arg.startswith(flag + "=") for flag in GLOBAL_FLAGS_WITH_VALUE if flag.startswith("--")):
            index += 1
            continue
        if arg.startswith("-"):
            index += 1
            continue
        return arg.lower()
    raise ValueError("无法识别 kubectl 操作类型")


READ_ONLY_CONFIG_SUBCOMMANDS = {"current-context", "get-clusters", "get-contexts", "get-users", "view"}


def classify_kubectl_command(command: str) -> CommandDecision:
    args = parse_kubectl_command(command)
    verb = _find_verb(args)
    if verb == "config":
        try:
            verb_index = args.index(next(arg for arg in args if arg.lower() == "config"))
            subcommand = args[verb_index + 1].lower()
        except (StopIteration, ValueError, IndexError):
            return CommandDecision("write", True, verb, args)
        if subcommand in READ_ONLY_CONFIG_SUBCOMMANDS:
            return CommandDecision("read", False, verb, args)
        return CommandDecision("write", True, verb, args)
    if verb in READ_ONLY_VERBS:
        return CommandDecision("read", False, verb, args)
    if verb in WRITE_OR_INTERACTIVE_VERBS:
        return CommandDecision("write", True, verb, args)
    return CommandDecision("write", True, verb, args)


# ── 系统诊断 shell 命令分类 (白名单) ─────────────────────

# 只读诊断命令白名单：value 为 None 表示该命令本身只读，任意参数均可；
# 为集合时，第一个非 flag 参数必须属于该集合。
SHELL_READ_ONLY_COMMANDS: dict[str, set[str] | None] = {
    "systemctl": {
        "status", "is-active", "is-enabled", "is-failed", "list-units",
        "list-sockets", "list-jobs", "show", "cat", "list-unit-files",
        "list-dependencies",
    },
    "journalctl": None,  # 只读，任意参数 (-u, -n, --no-pager, --since 等)
    "crictl": {"ps", "inspect", "logs", "info", "version", "images", "stats"},
    "ctr": {"containers", "images", "namespaces", "tasks", "snapshots", "info", "version", "plugins", "content", "leases"},
    "nerdctl": {"ps", "logs", "inspect", "info", "version", "images", "stats", "top", "events", "system", "namespace", "network", "volume"},
    "docker": {"ps", "logs", "inspect", "stats", "version", "info", "images", "top"},
    "df": None, "du": None, "free": None, "uptime": None, "uname": None, "hostname": None,
    "ip": {"addr", "address", "route", "link", "neighbor", "neigh", "rule", "maddr", "monitor", "tunnel", "tuntap"},
    "ss": None, "netstat": None, "lsblk": None, "ls": None, "cat": None,
    "head": None, "tail": None, "wc": None, "grep": None, "egrep": None, "fgrep": None,
    "ps": None, "pstree": None, "mount": None, "lsmod": None, "lscpu": None, "lsof": None,
    "dmesg": None, "ping": None, "nslookup": None, "dig": None, "getent": None,
    "timedatectl": None, "date": None, "stat": None, "file": None, "blkid": None,
    "top": None, "htop": None, "iostat": None, "vmstat": None, "mpstat": None,
    "pidstat": None, "sar": None, "nproc": None, "env": None, "which": None,
    "readlink": None, "realpath": None, "find": None, "sort": None, "uniq": None,
    "awk": None, "sed": None, "cut": None, "tr": None, "tee": None,
    "basename": None, "dirname": None, "xargs": None, "bc": None, "expr": None,
    "curl": None, "wget": None, "traceroute": None, "tracepath": None, "mtr": None,
    "nc": None, "nmap": None, "telnet": None, "tcpdump": None, "ethtool": None,
    "lspci": None, "lsusb": None, "dmidecode": None, "sysctl": None, "udevadm": None,
}

# etcdctl 由 classify_shell_command 中的专用块处理 (两段式命令)
# 从通用白名单中移除，避免被单段式逻辑误判

# etcdctl 单段式只读 verb
_ETCDCTL_READ_VERBS = {"get", "watch", "version", "lock", "lease", "move-leader", "leader"}
# etcdctl 单段式写 verb (需人工审批)
_ETCDCTL_WRITE_VERBS = {"put", "del", "delete", "compact", "txn", "snapshot", "defrag"}
# etcdctl 两段式 (verb, subverb) 只读组合
_ETCDCTL_READ_SUBVERBS = {
    "endpoint": {"status", "health", "hashkv"},
    "member": {"list", "list-urls"},
    "alarm": {"list"},
    "auth": {"status"},
    "user": {"list", "get"},
    "role": {"list", "get"},
    "check": {"perf", "datascale"},
}
# etcdctl 两段式 (verb, subverb) 写组合 (需人工审批)
_ETCDCTL_WRITE_SUBVERBS = {
    "member": {"add", "remove", "update"},
    "alarm": {"disarm"},
    "auth": {"enable", "disable"},
    "user": {"add", "delete", "grant-role", "revoke-role", "passwd"},
    "role": {"add", "delete", "grant-permission", "revoke-permission"},
}
# 系统修改命令白名单：需要人工审批
SHELL_WRITE_COMMANDS: dict[str, set[str] | None] = {
    "systemctl": {"restart", "start", "stop", "reload", "enable", "disable", "daemon-reload"},
    # 文件操作命令：任意参数都需审批
    "rm": None,
    "cp": None,
    "mv": None,
    "chmod": None,
    "chown": None,
    "mkdir": None,
    "touch": None,
    "ln": None,
}

# ip 命令的写操作动词 (出现即拒绝自动执行)
_IP_WRITE_ACTIONS = {"set", "add", "del", "delete", "change", "replace", "append", "prepend"}

# cat 允许读取的安全路径
_SAFE_CAT_FILES = {"/etc/os-release", "/etc/hosts", "/etc/resolv.conf"}
_SAFE_CAT_DIRS = ("/proc/", "/sys/", "/etc/kubernetes/", "/etc/cni/", "/etc/systemd/system/")


def _is_safe_cat_path(path: str) -> bool:
    if not path:
        return False
    if path in _SAFE_CAT_FILES:
        return True
    return any(path.startswith(d) for d in _SAFE_CAT_DIRS)


def _first_non_flag(args: list[str]) -> str | None:
    for arg in args:
        if not arg.startswith("-"):
            return arg.lower()
    return None


def parse_shell_command(command: str) -> list[str]:
    """解析单条 shell 诊断命令，禁止 shell 操作符和命令组合。"""
    parts = _parse_command(command)
    if not parts:
        raise ValueError("命令不能为空")
    return parts


def classify_shell_command(command: str) -> CommandDecision:
    """对白名单内的系统诊断命令进行分类。

    systemctl 同时出现在只读和写白名单中：restart/stop 等走写表（需审批），
    status/is-active 等走只读表（自动执行）。因此先检查写表，再检查只读表。

    含管道 (|) 或重定向 (>) 的命令：只检查管道前的 binary 是否在白名单内，
    并设置 needs_shell=True 以便用 shell 方式执行。
    """
    needs_shell = _has_shell_operators(command)

    if needs_shell:
        # 含 shell 操作符：排除分号（分号仍被禁止）
        if ';' in command:
            raise ValueError("不允许使用分号（;）连接多条命令")
        # 含管道/重定向：先检查管道后是否有危险命令
        if _pipe_to_dangerous_target(command):
            raise ValueError("管道后不允许执行危险命令（sh/rm/chmod 等），如需要请手动执行")
        # 取管道前的 binary 做白名单检查
        pipe_idx = command.index("|") if "|" in command else len(command)
        before_pipe = command[:pipe_idx].strip()
        try:
            before_parts = shlex.split(before_pipe)
        except ValueError:
            before_parts = [before_pipe]
        binary = before_parts[0]

        if binary not in SHELL_READ_ONLY_COMMANDS and binary not in SHELL_WRITE_COMMANDS:
            raise ValueError(f"不允许执行该命令: {binary}")
        # cat 的安全路径检查（管道场景）
        if binary == "cat":
            for arg in before_parts[1:]:
                if not arg.startswith("-") and not _is_safe_cat_path(arg):
                    raise ValueError(
                        "cat 仅允许读取 /proc, /sys, /etc/kubernetes, /etc/cni, "
                        "/etc/systemd/system, /etc/os-release, /etc/hosts, /etc/resolv.conf"
                    )
        # 对于含管道的命令，只检查 binary 白名单，不做子命令级别的检查
        # 同时拒绝写操作（如 systemctl restart kubelet | grep ... 也不允许自动执行）
        if binary in SHELL_WRITE_COMMANDS:
            allowed = SHELL_WRITE_COMMANDS[binary]
            if allowed is None:
                return CommandDecision("write", True, binary, [command], needs_shell=True)
            # 取管道前的子命令
            verb = _first_non_flag(before_parts[1:])
            if verb is not None and verb in allowed:
                return CommandDecision("write", True, binary, [command], needs_shell=True)
            # binary 在写表但子命令不匹配写表 — 走只读白名单检查
            if binary in SHELL_READ_ONLY_COMMANDS:
                ro_allowed = SHELL_READ_ONLY_COMMANDS[binary]
                if ro_allowed is None or (verb is not None and verb in ro_allowed):
                    return CommandDecision("read", False, binary, [command], needs_shell=True)
            raise ValueError(f"命令 {binary} 不支持该子命令: {verb}")

        return CommandDecision("read", False, binary, [command], needs_shell=True)

    parts = parse_shell_command(command)
    binary = parts[0]
    verb = _first_non_flag(parts[1:])

    if binary in SHELL_WRITE_COMMANDS:
        allowed = SHELL_WRITE_COMMANDS[binary]
        if allowed is None or (verb is not None and verb in allowed):
            return CommandDecision("write", True, binary, parts[1:])

    # etcdctl: 单段式或两段式命令，verb/subverb 均跳过全局 flag (--endpoints 等)
    # member/alarm/auth 等同时有只读和写子命令，需按 (verb, subverb) 组合精确判断
    if binary == "etcdctl":
        non_flags = [a for a in parts[1:] if not a.startswith("-")]
        ev = non_flags[0].lower() if non_flags else None
        esub = non_flags[1].lower() if len(non_flags) > 1 else None
        # 单段式 verb
        if ev in _ETCDCTL_READ_VERBS:
            return CommandDecision("read", False, ev, parts[1:])
        if ev in _ETCDCTL_WRITE_VERBS:
            return CommandDecision("write", True, ev, parts[1:])
        # 两段式: 先查写子命令 (精确匹配 subverb)，再查只读子命令
        if ev in _ETCDCTL_WRITE_SUBVERBS and esub in _ETCDCTL_WRITE_SUBVERBS[ev]:
            return CommandDecision("write", True, ev, parts[1:])
        if ev in _ETCDCTL_READ_SUBVERBS:
            if esub is None or esub in _ETCDCTL_READ_SUBVERBS[ev]:
                return CommandDecision("read", False, ev, parts[1:])
            raise ValueError(f"etcdctl {ev} 不支持子命令: {esub}")
        raise ValueError(f"etcdctl 不支持该子命令: {ev}")

    if binary in SHELL_READ_ONLY_COMMANDS:
        allowed = SHELL_READ_ONLY_COMMANDS[binary]
        if allowed is not None and (verb is None or verb not in allowed):
            raise ValueError(f"命令 {binary} 不支持该子命令: {verb}")
        if binary == "ip" and any(arg in _IP_WRITE_ACTIONS for arg in parts[1:]):
            raise ValueError("ip 命令包含写操作，请使用只读形式如 'ip addr show'")
        if binary == "cat":
            for arg in parts[1:]:
                if not arg.startswith("-") and not _is_safe_cat_path(arg):
                    raise ValueError(
                        "cat 仅允许读取 /proc, /sys, /etc/kubernetes, /etc/cni, "
                        "/etc/systemd/system, /etc/os-release, /etc/hosts, /etc/resolv.conf"
                    )
        return CommandDecision("read", False, binary, parts[1:])

    raise ValueError(f"不允许执行该命令: {binary}")


def classify_agent_command(command: str) -> CommandDecision:
    """统一分类 Agent 命令：kubectl 或白名单系统诊断命令。"""
    stripped = (command or "").strip()
    if stripped.startswith("kubectl ") or stripped == "kubectl":
        return classify_kubectl_command(stripped)
    return classify_shell_command(stripped)


# ── 命令执行 ──────────────────────────────────────────────

async def execute_command(command: str, timeout: int = 30) -> dict:
    """执行 Agent 命令 (kubectl 或白名单系统诊断命令)。

    含管道/重定向的命令通过 shell=True 执行，其他用 subprocess_exec 直接执行。
    """
    decision = classify_agent_command(command)
    parts = _parse_command(command)
    if parts[0] == "kubectl":
        executable = KUBECTL
        args = decision.args
    else:
        executable = parts[0]
        args = parts[1:]

    if decision.needs_shell:
        proc = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
    else:
        proc = await asyncio.create_subprocess_exec(
            executable,
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        return {
            "command": command,
            "mode": decision.mode,
            "returncode": proc.returncode,
            "stdout": stdout.decode(errors="replace")[:20000].strip(),
            "stderr": stderr.decode(errors="replace")[:10000].strip(),
        }
    except asyncio.TimeoutError:
        proc.kill()
        await proc.communicate()
        return {
            "command": command,
            "mode": decision.mode,
            "returncode": -1,
            "stdout": "",
            "stderr": f"命令执行超时 ({timeout}s)",
        }


async def execute_kubectl(command: str, timeout: int = 30) -> dict:
    """执行单条 kubectl 命令 (兼容旧接口)。"""
    decision = classify_kubectl_command(command)
    proc = await asyncio.create_subprocess_exec(
        KUBECTL,
        *decision.args,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        return {
            "command": command,
            "mode": decision.mode,
            "returncode": proc.returncode,
            "stdout": stdout.decode(errors="replace")[:20000].strip(),
            "stderr": stderr.decode(errors="replace")[:10000].strip(),
        }
    except asyncio.TimeoutError:
        proc.kill()
        await proc.communicate()
        return {
            "command": command,
            "mode": decision.mode,
            "returncode": -1,
            "stdout": "",
            "stderr": f"命令执行超时 ({timeout}s)",
        }