# 中文注释：
# 文件：echobot/agent/security.py
# 说明：安全策略、审批流与审计日志。

"""安全护栏模块：策略评估、人工审批与审计日志。"""

from __future__ import annotations

import hashlib
import json
import re
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Literal

from loguru import logger

from echobot.utils.helpers import ensure_dir, truncate_string

DecisionAction = Literal["allow", "approve", "deny"]
ApprovalStatus = Literal["pending", "approved", "rejected", "expired", "consumed"]

# 默认“需要审批”的工具集合。
# 说明：exec 与技能脚本执行都具备命令执行能力；文件写入/编辑会改变工作区状态。
# 因此这些能力默认都要求审批。
DEFAULT_APPROVAL_REQUIRED_TOOLS: set[str] = {
    "exec",
    "write_file",
    "edit_file",
    "run_skill_script",
}

# 默认“exec 直接拒绝”的命令模式。
# 说明：命中这些模式会直接 deny，不进入审批队列，避免误放行破坏性命令。
DEFAULT_DENIED_EXEC_PATTERNS: list[str] = [
    r"(^|[;&|])\s*rm\s+-rf\s+/\s*($|[;&|])",
    r"(^|[;&|])\s*shutdown\b",
    r"(^|[;&|])\s*reboot\b",
    r"(^|[;&|])\s*mkfs(\.| )",
    r"(^|[;&|])\s*dd\s+if=",
    r"curl\s+[^|]*\|\s*(sh|bash)",
    r"wget\s+[^|]*\|\s*(sh|bash)",
]


def _utc_now() -> datetime:
    """返回 UTC 当前时间（统一时区，避免多机/夏令时差异）。"""
    return datetime.now(timezone.utc)


@dataclass
class SecurityDecision:
    """
    安全策略评估结果。

    action:
    - allow: 直接放行
    - approve: 需人工审批
    - deny: 直接拒绝
    """

    action: DecisionAction
    reason: str = ""


@dataclass
class SecurityGateResult:
    """
    运行时门禁结果。

    该结构用于返回给 AgentLoop：
    - 是否可执行
    - 给模型/用户的反馈文案
    - 如有审批请求，带 request_id
    """

    action: DecisionAction
    message: str
    request_id: str | None = None


@dataclass
class ApprovalRequest:
    """
    一条审批请求记录（待处理或已处理）。

    signature 用于“同一会话 + 同一工具 + 同一参数”的去重与匹配。
    """

    id: str
    session_key: str
    tool_name: str
    params: dict[str, Any]
    reason: str
    signature: str
    status: ApprovalStatus = "pending"
    created_at: str = field(default_factory=lambda: _utc_now().isoformat())
    expires_at: str = ""
    decided_at: str | None = None
    decided_by: str | None = None
    consumed_at: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


class SecurityPolicy:
    """
    工具调用安全策略（简版）。

    设计目标：
    1. 先覆盖高风险能力（shell、文件写入）；
    2. 给出可解释的拒绝原因；
    3. 后续可平滑扩展为基于角色/租户/路径白名单的细粒度策略。
    """

    def __init__(
        self,
        approval_required_tools: set[str] | None = None,
        denied_exec_patterns: list[str] | None = None,
    ):
        # 中文注释：
        # 这里对配置输入做“规范化”处理，保证后续评估逻辑只处理干净数据：
        # - tool 名称去空白、去空字符串
        # - 正则模式提前校验，过滤非法模式，避免运行时抛异常
        self.approval_required_tools = self._normalize_approval_required_tools(
            approval_required_tools
        )
        self.denied_exec_patterns = self._normalize_denied_exec_patterns(denied_exec_patterns)

    @staticmethod
    def _normalize_approval_required_tools(tools: set[str] | None) -> set[str]:
        """
        归一化“需要审批”的工具集合。

        语义约定：
        - None: 使用默认值
        - 空集合 set(): 显式表示“无工具需要审批”
        """
        if tools is None:
            return set(DEFAULT_APPROVAL_REQUIRED_TOOLS)

        normalized: set[str] = set()
        for name in tools:
            value = str(name).strip()
            if value:
                normalized.add(value)
        return normalized

    @staticmethod
    def _normalize_denied_exec_patterns(patterns: list[str] | None) -> list[str]:
        """
        归一化并校验 exec 拒绝规则。

        语义约定：
        - None: 使用默认值
        - 空列表 []: 显式表示不启用“直接拒绝模式”
        """
        source = DEFAULT_DENIED_EXEC_PATTERNS if patterns is None else patterns
        normalized: list[str] = []
        for raw in source:
            pattern = str(raw).strip()
            if not pattern:
                continue
            try:
                # 预编译用于校验 regex 合法性，失败则跳过该条规则。
                re.compile(pattern, re.IGNORECASE)
            except re.error as e:
                logger.warning(f"Skip invalid denied_exec_pattern '{pattern}': {e}")
                continue
            normalized.append(pattern)
        return normalized

    def evaluate(self, tool_name: str, params: dict[str, Any]) -> SecurityDecision:
        """
        对单次工具调用做策略评估。

        当前规则：
        - exec: 先做危险模式拦截；再按配置决定 approve/allow
        - 其他工具：若在 approval_required_tools 中则审批
        - 其他工具: 直接放行
        """
        if tool_name == "exec":
            command = str(params.get("command", ""))
            for pattern in self.denied_exec_patterns:
                if re.search(pattern, command, re.IGNORECASE):
                    # 命中阻断模式直接拒绝，避免进入审批后误放行。
                    return SecurityDecision(
                        action="deny",
                        reason=f"Command matches blocked pattern: {pattern}",
                    )
            # exec 是否审批由配置控制：
            # - 包含 exec -> approve
            # - 不包含 exec -> allow
            if "exec" in self.approval_required_tools:
                return SecurityDecision(
                    action="approve",
                    reason="Shell commands require human approval.",
                )
            return SecurityDecision(action="allow", reason="Shell command allowed by policy.")

        if tool_name in self.approval_required_tools:
            return SecurityDecision(action="approve", reason=f"Tool '{tool_name}' requires approval.")

        return SecurityDecision(action="allow", reason="Allowed by policy.")


class ApprovalManager:
    """
    审批请求生命周期管理器。

    能力：
    - 创建/去重 pending 请求
    - 审批通过/拒绝
    - 过期清理
    - 消费一次性批准令牌
    """

    def __init__(self, ttl_seconds: int = 15 * 60, max_pending: int = 200):
        self.ttl_seconds = max(60, ttl_seconds)
        self.max_pending = max(10, max_pending)
        self._requests: dict[str, ApprovalRequest] = {}

    @staticmethod
    def _signature(session_key: str, tool_name: str, params: dict[str, Any]) -> str:
        """为审批请求生成稳定签名，用于幂等匹配。"""
        payload = json.dumps(
            {"session_key": session_key, "tool_name": tool_name, "params": params},
            ensure_ascii=False,
            sort_keys=True,
            default=str,
        )
        return hashlib.sha256(payload.encode("utf-8")).hexdigest()

    @staticmethod
    def _parse_iso(ts: str) -> datetime:
        """解析 ISO 时间字符串。"""
        return datetime.fromisoformat(ts)

    def _cleanup_expired(self) -> None:
        """将超时的 pending/approved 请求标记为 expired。"""
        now = _utc_now()
        for req in self._requests.values():
            if req.status not in {"pending", "approved"}:
                continue
            try:
                if self._parse_iso(req.expires_at) < now:
                    req.status = "expired"
            except Exception:
                req.status = "expired"

    def _trim_pending_if_needed(self) -> None:
        """当 pending 数量过多时，淘汰最老请求，防止内存无限增长。"""
        pending = [r for r in self._requests.values() if r.status == "pending"]
        if len(pending) < self.max_pending:
            return

        pending.sort(key=lambda r: r.created_at)
        overflow = len(pending) - self.max_pending + 1
        for req in pending[:overflow]:
            req.status = "expired"

    def create_request(
        self,
        session_key: str,
        tool_name: str,
        params: dict[str, Any],
        reason: str,
    ) -> ApprovalRequest:
        """
        创建审批请求，若存在同签名 pending 则复用。

        这样可以避免模型在短时间内重复发起同一审批弹窗。
        """
        self._cleanup_expired()
        self._trim_pending_if_needed()
        signature = self._signature(session_key, tool_name, params)

        for req in self._requests.values():
            if req.status == "pending" and req.signature == signature:
                return req

        req = ApprovalRequest(
            id=f"apr_{uuid.uuid4().hex[:12]}",
            session_key=session_key,
            tool_name=tool_name,
            params=params,
            reason=reason,
            signature=signature,
            expires_at=(_utc_now() + timedelta(seconds=self.ttl_seconds)).isoformat(),
        )
        self._requests[req.id] = req
        return req

    def list_pending(self) -> list[dict[str, Any]]:
        """返回全部待审批请求（按创建时间倒序）。"""
        self._cleanup_expired()
        pending = [r.to_dict() for r in self._requests.values() if r.status == "pending"]
        pending.sort(key=lambda x: x["created_at"], reverse=True)
        return pending

    def get_request(self, request_id: str) -> dict[str, Any] | None:
        """按 request_id 查询审批请求。"""
        req = self._requests.get(request_id)
        return req.to_dict() if req else None

    def approve(self, request_id: str, approver: str = "user") -> bool:
        """审批通过。仅 pending 状态可转为 approved。"""
        self._cleanup_expired()
        req = self._requests.get(request_id)
        if not req or req.status != "pending":
            return False
        req.status = "approved"
        req.decided_by = approver
        req.decided_at = _utc_now().isoformat()
        return True

    def reject(self, request_id: str, approver: str = "user") -> bool:
        """审批拒绝。仅 pending 状态可转为 rejected。"""
        self._cleanup_expired()
        req = self._requests.get(request_id)
        if not req or req.status != "pending":
            return False
        req.status = "rejected"
        req.decided_by = approver
        req.decided_at = _utc_now().isoformat()
        return True

    def consume_approved(
        self, session_key: str, tool_name: str, params: dict[str, Any]
    ) -> str | None:
        """
        消费“已批准”请求（一次性）。

        返回值：
        - 匹配到可消费批准 -> request_id
        - 否则 -> None

        这样可避免“批准一次，永久放行”。
        """
        self._cleanup_expired()
        signature = self._signature(session_key, tool_name, params)

        approved = [
            r
            for r in self._requests.values()
            if r.status == "approved" and r.signature == signature
        ]
        if not approved:
            return None

        approved.sort(key=lambda r: r.decided_at or "")
        req = approved[-1]
        req.status = "consumed"
        req.consumed_at = _utc_now().isoformat()
        return req.id


class SecurityAuditLogger:
    """
    审计日志写入器（JSONL）。

    记录每次策略决策与审批动作，便于：
    - 问题排查
    - 安全复盘
    - 合规追踪
    """

    def __init__(self, workspace: Path):
        # 审计日志放在工作区内，便于实例隔离和打包留存。
        self.log_path = ensure_dir(workspace / "security") / "audit.log"

    def log(self, event: str, **payload: Any) -> None:
        """写入单条审计事件（失败仅告警，不阻塞主流程）。"""
        data = {"timestamp": _utc_now().isoformat(), "event": event, **payload}
        try:
            with open(self.log_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(data, ensure_ascii=False, default=str) + "\n")
        except Exception as e:
            logger.warning(f"Failed to write security audit log: {e}")


class SecurityManager:
    """
    安全管理门面。

    对外提供统一入口：
    - 策略评估
    - 审批管理
    - 审计落盘
    """

    def __init__(
        self,
        workspace: Path,
        approval_ttl_seconds: int = 15 * 60,
        approval_required_tools: set[str] | None = None,
        denied_exec_patterns: list[str] | None = None,
    ):
        # 采用组合方式，后续可无缝替换策略引擎/审批存储后端。
        self.policy = SecurityPolicy(
            approval_required_tools=approval_required_tools,
            denied_exec_patterns=denied_exec_patterns,
        )
        self.approvals = ApprovalManager(ttl_seconds=approval_ttl_seconds)
        self.audit = SecurityAuditLogger(workspace=workspace)

    def check_tool_execution(
        self, session_key: str, tool_name: str, params: dict[str, Any]
    ) -> SecurityGateResult:
        """
        对一次工具执行进行门禁判定。

        执行顺序：
        1. 策略评估（allow/approve/deny）
        2. 若 approve，尝试消费已批准令牌
        3. 若无可消费令牌，创建 pending 请求
        4. 全程写审计日志
        """
        decision = self.policy.evaluate(tool_name=tool_name, params=params)

        # 仅记录参数摘要，避免审计日志写入超长敏感正文。
        params_preview = truncate_string(json.dumps(params, ensure_ascii=False, default=str), 400)

        if decision.action == "deny":
            self.audit.log(
                "tool.execution.denied",
                session_key=session_key,
                tool_name=tool_name,
                reason=decision.reason,
                params_preview=params_preview,
            )
            return SecurityGateResult(action="deny", message=decision.reason)

        if decision.action == "allow":
            self.audit.log(
                "tool.execution.allowed",
                session_key=session_key,
                tool_name=tool_name,
                params_preview=params_preview,
            )
            return SecurityGateResult(action="allow", message=decision.reason)

        approved_request_id = self.approvals.consume_approved(
            session_key=session_key, tool_name=tool_name, params=params
        )
        if approved_request_id:
            self.audit.log(
                "tool.execution.approval_consumed",
                request_id=approved_request_id,
                session_key=session_key,
                tool_name=tool_name,
                params_preview=params_preview,
            )
            return SecurityGateResult(
                action="allow",
                message="Approval matched and consumed.",
                request_id=approved_request_id,
            )

        req = self.approvals.create_request(
            session_key=session_key,
            tool_name=tool_name,
            params=params,
            reason=decision.reason,
        )
        self.audit.log(
            "tool.execution.approval_requested",
            request_id=req.id,
            session_key=session_key,
            tool_name=tool_name,
            reason=req.reason,
            params_preview=params_preview,
        )
        return SecurityGateResult(
            action="approve",
            message=(
                f"Approval required for tool '{tool_name}'. "
                f"request_id={req.id}. Approve then retry the request."
            ),
            request_id=req.id,
        )

    def list_pending_approvals(self) -> list[dict[str, Any]]:
        """列出所有待审批请求。"""
        return self.approvals.list_pending()

    def approve(self, request_id: str, approver: str = "user") -> bool:
        """批准请求，并记录审计日志。"""
        ok = self.approvals.approve(request_id, approver=approver)
        if ok:
            self.audit.log(
                "tool.execution.approval_decided",
                request_id=request_id,
                decision="approved",
                decided_by=approver,
            )
        return ok

    def reject(self, request_id: str, approver: str = "user") -> bool:
        """拒绝请求，并记录审计日志。"""
        ok = self.approvals.reject(request_id, approver=approver)
        if ok:
            self.audit.log(
                "tool.execution.approval_decided",
                request_id=request_id,
                decision="rejected",
                decided_by=approver,
            )
        return ok
