"""Approval workflow for the LangGraph agent.

Implements a 3-level policy engine (ALLOW / DENY / ASK_USER) with 4 approval
modes and session-scoped approval caching. Ported from Gemini CLI's
policy-engine.ts and Codex CLI's exec_policy.rs.

Key design decisions:
- Pre-LLM tool filtering: in PLAN mode, write tools are removed entirely
  so the LLM never wastes tokens trying to call them.
- Approval caching: when user says "always", cache for the session.
- Command parsing for bash: split on shell operators, evaluate each sub-command.
"""

from __future__ import annotations

import logging
import re
import time
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


class ApprovalMode(str, Enum):
    """User-selectable approval modes."""
    DEFAULT = "default"        # Ask for writes + dangerous commands
    AUTO_EDIT = "auto-edit"    # Auto-approve file edits, ask for commands
    YOLO = "yolo"              # Auto-approve everything (current behavior)
    PLAN = "plan"              # Read-only tools only (writes/bash denied)


class PolicyDecision(str, Enum):
    """Result of evaluating a policy rule."""
    ALLOW = "allow"
    DENY = "deny"
    ASK_USER = "ask_user"


# Tools that are always safe (read-only)
READ_ONLY_TOOLS = frozenset({
    "read_file", "search_files", "list_directory",
    "git_log", "git_diff", "git_blame",
    "parse_error", "ask_user", "submit_plan",
})

# Shell injection patterns — always dangerous regardless of command prefix.
# Catches command substitution, parameter expansion, process substitution,
# and execution wrappers that can hide arbitrary commands inside "safe" prefixes.
_SHELL_INJECTION_RE = re.compile(
    r"\$\("          # Command substitution: $(...)
    r"|`"            # Backtick substitution: `...`
    r"|\$\{"         # Parameter expansion: ${...}
    r"|<\(|>\("      # Process substitution: <(...) or >(...)
    r"|\beval\b"     # eval command
    r"|(?:^|\s)exec\b" # exec command (not -exec flag)
    r"|\bsource\b"   # source command
    r"|\balias\b"    # alias can redefine commands
    r"|\$'[^']*'"    # ANSI-C quoting: $'\x72\x6d' can hide commands
    # NOTE: dot-source (`. script.sh`) checked per-sub-command in _is_dot_source(),
    # not here, because `(?:^|\s)\.\s` false-positives on `find . -name '*.py'`.
)

# Shell separators we don't parse — presence means "not safe, needs approval".
# The command parser (_parse_bash_commands) handles &&, ||, ;, |. Anything else
# that separates commands in bash is an unhandled separator. Instead of trying
# to split on every possible separator (whack-a-mole), we detect them and
# flag the whole command as not-auto-approvable.
_UNHANDLED_SHELL_RE = re.compile(
    r"(?<![&])\&(?![&>])"  # Bare & (background exec), not && or &>
    r"|\n"                  # Newline (command separator in bash)
)

# Exec-embedding flags — commands that embed arbitrary execution in their arguments.
# Unlike python -c / node -e (removed from safe list because EVERY usage is code
# execution), these are commands with both safe and dangerous usages. We detect
# the dangerous flags specifically rather than removing the command entirely.
_EXEC_EMBEDDING_RE = re.compile(
    r"(?:^|\s)-exec\s"          # find -exec (followed by space, not -executable)
    r"|(?:^|\s)-execdir\s"      # find -execdir
    r"|(?:^|\s)-ok\s"           # find -ok (like -exec but prompts)
    r"|(?:^|\s)-delete(?:\s|$)" # find -delete (deletes every matched file)
)

# Output redirect patterns — not injection, but file write operations that shouldn't auto-approve.
# Catches >, >>, n>, n>> (where n is a file descriptor digit) but NOT:
# - process substitution >( which is already caught by _SHELL_INJECTION_RE
# - &> and &>> (also caught here as write operations)
_OUTPUT_REDIRECT_RE = re.compile(
    r"(?<!\()"       # Negative lookbehind: not preceded by ( (skip process substitution >()
    r"\d*>{1,2}"     # Optional fd digit + > or >>
    r"(?!\()"        # Negative lookahead: not followed by ( (skip process substitution >()
)

# Bash commands that are always safe (read-only)
_SAFE_BASH_PREFIXES = frozenset({
    "ls", "cat", "head", "tail", "wc", "grep", "rg", "find", "which", "file",
    "echo", "printf", "pwd", "printenv", "whoami", "hostname", "uname", "date",
    # NOTE: "env" intentionally excluded — `env rm -rf /` runs rm
    "git status", "git log", "git diff", "git show", "git blame", "git branch",
    "git tag", "git remote", "git stash list",
    "python -m pytest", "python -m unittest",
    "pytest",
    "npm test", "npx jest", "yarn test", "cargo test", "go test",
    # NOTE: "python -c" and "node -e" intentionally excluded — interpreter
    # wrappers that can execute arbitrary code (same threat class as "env")
    "mypy", "ruff", "eslint", "tsc --noEmit",
    "pip list", "pip show", "npm list", "cargo --version",
})

# Transparent command wrappers — they run the command unchanged,
# only modifying scheduling/priority/session/tracing behavior.
# Used by _strip_transparent_wrappers() in _is_dangerous_bash_command
# to reveal the inner command for danger classification.
_TRANSPARENT_WRAPPERS = frozenset({
    "time", "timeout", "nice", "nohup", "strace", "ltrace",
    "ionice", "taskset", "chrt", "setsid",
    "xargs",  # xargs passes args to inner command (adversarial S9)
})

# Bash commands that always need approval
_DANGEROUS_BASH_PREFIXES = frozenset({
    "rm",
    "git push", "git reset", "git checkout", "git rebase",
    "git commit", "git clean",
    "git stash drop", "git stash clear", "git stash pop",
    "git stash apply", "git stash",
    "docker", "kubectl", "terraform",
    "curl", "wget", "ssh", "scp", "rsync",
    "pip install", "pip uninstall", "npm install", "npm uninstall",
    "yarn add", "yarn remove", "cargo install",
    "chmod", "chown", "mv", "cp", "kill", "pkill", "killall",
    "sudo", "env",
    # Destructive disk/device tools (fuzz finding M-1, adversarial S11)
    "mkfs", "mke2fs", "mkswap", "dd",
    "fdisk", "parted", "shred", "wipefs", "badblocks",
    # Schedulers can defer dangerous commands (adversarial S6)
    "at", "crontab",
})

# Pre-sorted by length (longest first) so specific patterns match before general.
_DANGEROUS_BASH_SORTED = sorted(_DANGEROUS_BASH_PREFIXES, key=len, reverse=True)


class ApprovalCache:
    """Session-scoped cache for approved tool calls.

    Supports three scopes:
    - exact: cache this exact tool+args combination
    - prefix: cache all commands with this prefix (e.g., "pytest" covers "pytest tests/")
    - tool: cache all calls to this tool
    """

    def __init__(self) -> None:
        self._exact: dict[str, float] = {}
        self._prefixes: dict[str, float] = {}   # tool:prefix → timestamp
        self._tools: dict[str, float] = {}       # tool_name → timestamp

    def is_approved(self, tool_name: str, args: dict) -> bool:
        """Check if this tool call is pre-approved."""
        # Tool-level approval
        if tool_name in self._tools:
            return True # The prefix in _prefixes already includes "bash:"

        # Prefix-level approval (for bash commands)
        if tool_name == "bash":
            cmd = args.get("command", "").strip()
            for prefix in self._prefixes:
                if prefix.startswith("bash:") and _matches_prefix_with_word_boundary(cmd, prefix[5:]):
                    return True

        # Exact approval
        import json
        key = f"{tool_name}:{json.dumps(args, sort_keys=True, default=str)}"
        return key in self._exact

    def approve(self, tool_name: str, args: dict, scope: str = "exact") -> None:
        """Add an approval to the cache.

        Args:
            tool_name: Tool name.
            args: Tool arguments.
            scope: "exact" | "prefix" | "tool"
        """
        now = time.time()
        if scope == "tool":
            self._tools[tool_name] = now
        elif scope == "prefix" and tool_name == "bash":
            cmd = args.get("command", "").strip()
            # Extract command prefix (first word or first two words)
            parts = cmd.split()
            prefix = parts[0] if parts else cmd
            if len(parts) > 1 and parts[0] in ("python", "npm", "npx", "cargo", "git"):
                prefix = f"{parts[0]} {parts[1]}"
            self._prefixes[f"bash:{prefix}"] = now
        else:
            import json
            key = f"{tool_name}:{json.dumps(args, sort_keys=True, default=str)}"
            self._exact[key] = now


def _parse_bash_commands(command: str) -> list[str]:
    """Split a bash command into sub-commands.

    Splits on &&, ||, ;, and | to evaluate each independently.
    """
    # Split on shell operators but keep it simple (no nested quotes handling)
    parts = re.split(r'\s*(?:&&|\|\||;|\|)\s*', command.strip())
    return [p.strip() for p in parts if p.strip()]


def _matches_prefix_with_word_boundary(command: str, prefix: str) -> bool:
    """Check if command matches prefix with word boundary.
    
    Matches if:
    - command exactly equals prefix
    - command starts with prefix + space/tab/newline
    
    This prevents false matches like:
    - 'lsblk' matching 'ls'
    - 'echooooo' matching 'echo'
    - 'catastrophe' matching 'cat'
    """
    if command == prefix:
        return True
    if command.startswith(prefix + " "):
        return True
    if command.startswith(prefix + "\t"):
        return True
    if command.startswith(prefix + "\n"):
        return True
    return False


def _is_dot_source(sub_command: str) -> bool:
    """Check if a sub-command is dot-source (`. script.sh`).

    Dot-source is a command, not an argument, so it only appears at the
    start of a sub-command. Checked per-sub-command instead of via regex
    on the full command string, because `(?:^|\\s)\\.\\s` false-positives
    on commands like `find . -name '*.py'` where `.` is a directory arg.
    """
    return sub_command == "." or sub_command.startswith(". ")


def _strip_transparent_wrappers(sub_command: str) -> list[str]:
    """Return candidate inner commands after stripping transparent wrappers.

    Used by _is_dangerous_bash_command to classify wrapped commands.
    `time rm -rf /` → candidates include `rm -rf /` (dangerous)
    `nice -n 19 rm -rf /` → candidates include `rm -rf /` (dangerous)
    `time nice nohup rm -rf /` → candidates include `rm -rf /` (recursive)

    NOT used by _is_safe_bash_command — `time ls` should still route to
    ASK_USER because the wrapper itself is not in the safe list. But the
    approval prompt should show DANGEROUS (not generic "unknown command")
    because the inner command IS dangerous.

    Strategy: instead of trying to parse each wrapper's flag syntax (fragile —
    nice -n 19, timeout -s KILL 30, strace -e trace=open all differ), we
    try every suffix of the token list after the wrapper keyword. The caller
    checks each candidate against the dangerous prefix list. False negatives
    just mean the command stays UNSAFE (ASK_USER) instead of DANGEROUS.
    """
    parts = sub_command.split()
    if not parts or parts[0] not in _TRANSPARENT_WRAPPERS:
        return []
    # Skip past all leading wrapper keywords (handles time nice nohup ...)
    i = 0
    while i < len(parts) and parts[i] in _TRANSPARENT_WRAPPERS:
        i += 1
    # Every remaining suffix is a candidate inner command
    candidates = []
    for j in range(i, len(parts)):
        candidates.append(" ".join(parts[j:]))
    return candidates


def _is_safe_bash_command(command: str) -> bool:
    """Check if a bash command is safe to auto-approve."""
    # Shell injection = never safe, regardless of prefix
    if _SHELL_INJECTION_RE.search(command):
        return False
    # Output redirects = file write operations, not auto-approvable
    if _OUTPUT_REDIRECT_RE.search(command):
        return False
    # Unhandled shell separators = we can't parse these, needs approval
    if _UNHANDLED_SHELL_RE.search(command):
        return False
    # Exec-embedding flags = commands that embed arbitrary execution in args
    if _EXEC_EMBEDDING_RE.search(command):
        return False
    sub_commands = _parse_bash_commands(command)
    if not sub_commands:
        return False
    for sub in sub_commands:
        # Dot-source (`. script.sh`) is injection — checked per-sub-command
        if _is_dot_source(sub):
            return False
        # Variable indirection in command position (S2): if the first token
        # of a sub-command starts with `$` followed by a word character, the
        # shell will execute whatever value that variable holds. This is
        # suspicious (`cmd=rm; $cmd -rf /`) but not necessarily dangerous,
        # so we route to ASK_USER by returning not-safe.
        # We exclude `$(` which is already caught by _SHELL_INJECTION_RE.
        first_token = sub.split()[0] if sub.split() else ""
        if first_token and re.match(r"^\$\w", first_token):
            return False
        is_safe = any(
            _matches_prefix_with_word_boundary(sub, prefix)
            for prefix in _SAFE_BASH_PREFIXES
        )
        if not is_safe:
            return False
    return True


def _is_dangerous_bash_command(command: str) -> bool:
    """Check if a bash command contains dangerous operations.

    NOTE: The caller (PolicyEngine.check) calls _is_safe_bash_command first,
    so safe commands never reach here in normal flow. However, this function
    may be called directly (tests, fuzzer), so we guard against prefix
    conflicts where a command matches both a safe prefix (longer) and a
    dangerous prefix (shorter), e.g. "git stash list" matches safe
    "git stash list" and dangerous "git stash".
    """
    # Shell injection = always dangerous
    if _SHELL_INJECTION_RE.search(command):
        return True
    sub_commands = _parse_bash_commands(command)
    if not sub_commands:
        return False
    for sub in sub_commands:
        # Dot-source is injection
        if _is_dot_source(sub):
            return True
        # If this sub-command matches a safe prefix, skip dangerous check
        # (safe is a longer, more specific match than the dangerous prefix)
        sub_is_safe = any(
            _matches_prefix_with_word_boundary(sub, sp)
            for sp in _SAFE_BASH_PREFIXES
        )
        if sub_is_safe:
            continue
        # Check the sub-command as-is
        for prefix in _DANGEROUS_BASH_SORTED:
            if _matches_prefix_with_word_boundary(sub, prefix):
                return True
        # Strip transparent wrappers and check inner command candidates
        for candidate in _strip_transparent_wrappers(sub):
            for prefix in _DANGEROUS_BASH_SORTED:
                if _matches_prefix_with_word_boundary(candidate, prefix):
                    return True

    return False


class PolicyEngine:
    """Policy engine for tool approval.

    Evaluates tool calls against the current approval mode and cache.
    """

    def __init__(
        self,
        mode: ApprovalMode = ApprovalMode.YOLO,
        cache: Optional[ApprovalCache] = None,
    ) -> None:
        self.mode = mode
        self.cache = cache or ApprovalCache()

    def check(self, tool_name: str, args: dict) -> PolicyDecision:
        """Evaluate whether a tool call should be allowed.

        Args:
            tool_name: Name of the tool.
            args: Tool arguments.

        Returns:
            PolicyDecision: ALLOW, DENY, or ASK_USER.
        """
        # YOLO mode: everything allowed
        if self.mode == ApprovalMode.YOLO:
            return PolicyDecision.ALLOW

        # PLAN mode: only read-only tools
        if self.mode == ApprovalMode.PLAN:
            if tool_name in READ_ONLY_TOOLS:
                return PolicyDecision.ALLOW
            return PolicyDecision.DENY

        # Read-only tools are always allowed
        if tool_name in READ_ONLY_TOOLS:
            return PolicyDecision.ALLOW

        # Check cache
        if self.cache.is_approved(tool_name, args):
            return PolicyDecision.ALLOW

        # AUTO_EDIT mode: auto-approve file edits, ask for bash
        if self.mode == ApprovalMode.AUTO_EDIT:
            if tool_name in ("write_file", "edit_file"):
                return PolicyDecision.ALLOW
            if tool_name == "bash":
                cmd = args.get("command", "")
                if _is_safe_bash_command(cmd):
                    return PolicyDecision.ALLOW
                return PolicyDecision.ASK_USER
            # Other tools: allow
            return PolicyDecision.ALLOW

        # DEFAULT mode: ask for writes and dangerous bash
        if tool_name in ("write_file", "edit_file"):
            return PolicyDecision.ASK_USER

        if tool_name == "bash":
            cmd = args.get("command", "")
            if _is_safe_bash_command(cmd):
                return PolicyDecision.ALLOW
            return PolicyDecision.ASK_USER

        # run_tests uses bash internally — allow it
        if tool_name == "run_tests":
            return PolicyDecision.ALLOW

        # web_fetch: network access requires approval in DEFAULT mode
        if tool_name == "web_fetch":
            return PolicyDecision.ASK_USER

        # save_memory: writes to disk, requires approval
        if tool_name == "save_memory":
            return PolicyDecision.ASK_USER

        # Default: allow
        return PolicyDecision.ALLOW

    def get_available_tools(self, all_tools: list) -> list:
        """Filter tools based on current mode (pre-LLM filtering).

        In PLAN mode, removes write tools so the LLM never tries to call them.
        """
        if self.mode != ApprovalMode.PLAN:
            return all_tools
        return [t for t in all_tools if t.name in READ_ONLY_TOOLS]

    def format_approval_request(self, tool_name: str, args: dict) -> str:
        """Format an approval request for display."""
        if tool_name == "bash":
            cmd = args.get("command", "")
            return f"bash: {cmd}"
        elif tool_name in ("write_file", "edit_file"):
            path = args.get("path", args.get("file_path", "unknown"))
            return f"{tool_name}: {path}"
        else:
            return f"{tool_name}: {args}"
