"""Shell execution toolkit with accuracy-oriented result enrichment.

Key design principles (inspired by OpenClaw):
  1. Rich result metadata --exit code classification, duration, cwd, failure reasons
  2. Preflight validation --script variable leakage, obfuscation, inline eval detection
  3. Output sanitization --binary character filtering, smart truncation
  4. Platform-aware execution --Windows PowerShell preference, process tree cleanup
"""

import asyncio
import os
import platform
import re
import shutil
import time
from pathlib import Path
from typing import Any

from loguru import logger

from agno.tools import Toolkit


# ── Exit code classification ────────────────────────────────────────

class ExitCodeInfo:
    """Classify exit codes into human-readable failure reasons."""

    SHELL_NOT_FOUND = 127       # command not found (Unix)
    SHELL_NOT_EXECUTABLE = 126  # permission denied (Unix)
    TIMEOUT_CODE = 124          # timeout (conventional)
    SIGKILL = 137               # 128 + 9
    SIGTERM = 143               # 128 + 15
    SIGINT = 130                # 128 + 2
    WIN_NOT_FOUND = 9009        # command not found (Windows cmd.exe)
    WIN_NOT_FOUND_PS = 1        # PowerShell command not found (generic)

    @classmethod
    def classify(cls, code: int | None, timed_out: bool = False) -> str | None:
        """Return a human-readable failure reason, or None if success."""
        if code is None:
            return "Process exited without status code"
        if code == 0:
            return None
        if timed_out or code == cls.TIMEOUT_CODE:
            return "Command timed out"
        if code in (cls.SHELL_NOT_FOUND, cls.WIN_NOT_FOUND):
            return "Command not found -- check spelling or install the required tool"
        if code == cls.SHELL_NOT_EXECUTABLE:
            return "Command not executable -- check file permissions"
        if code == cls.SIGKILL:
            return "Process killed (SIGKILL) --possibly OOM or force-killed"
        if code == cls.SIGTERM:
            return "Process terminated (SIGTERM)"
        if code == cls.SIGINT:
            return "Process interrupted (SIGINT)"
        if code > 128:
            sig = code - 128
            return f"Process killed by signal {sig}"
        return f"Command failed with exit code {code}"


# ── Output sanitization ─────────────────────────────────────────────

# Control characters to strip (keep tab \x09, newline \x0a, carriage return \x0d)
_CONTROL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")


def _decode_output(data: bytes) -> str:
    """Decode subprocess output, trying UTF-8 first then system encoding."""
    try:
        return data.decode("utf-8")
    except UnicodeDecodeError:
        pass
    # On Windows, console output is often in the system ANSI code page (e.g., GBK)
    if platform.system() == "Windows":
        for enc in ("gbk", "cp936", "cp1252"):
            try:
                return data.decode(enc)
            except (UnicodeDecodeError, LookupError):
                continue
    return data.decode("utf-8", errors="replace")


def sanitize_output(text: str) -> str:
    """Remove binary artifacts and control characters from command output."""
    # Strip control chars (keep \t \n \r)
    text = _CONTROL_RE.sub("", text)
    # Remove NUL bytes that survive decoding
    text = text.replace("\x00", "")
    # Remove Unicode format characters (zero-width spaces, direction markers, etc.)
    import unicodedata
    text = "".join(c for c in text if unicodedata.category(c) not in ("Cf", "Cs"))
    return text


def truncate_output(text: str, max_chars: int, strategy: str = "tail") -> tuple[str, bool]:
    """Truncate output with configurable strategy.

    Strategies:
      - "tail": keep last max_chars (most recent output, best for logs)
      - "head": keep first max_chars (traditional)
      - "middle": keep head + tail, cut middle (best for large structured output)

    Returns (truncated_text, was_truncated).
    """
    if len(text) <= max_chars:
        return text, False

    omitted = len(text) - max_chars
    marker = f"\n\n... ({omitted:,} characters omitted) ...\n\n"

    if strategy == "tail":
        return marker + text[-max_chars:], True
    elif strategy == "middle":
        half = (max_chars - len(marker)) // 2
        return text[:half] + marker + text[-half:], True
    else:  # head
        return text[:max_chars] + f"\n... ({omitted:,} more characters)", True


# ── Preflight validation ─────────────────────────────────────────────

# Shell variable syntax that shouldn't appear in Python/JS files
_SHELL_VAR_RE = re.compile(r"\$[A-Z_][A-Z0-9_]{1,}")

# Obfuscation patterns --encoded command execution
_OBFUSCATION_PATTERNS = [
    re.compile(r"base64\s+(-d|--decode)\s*\|.*\b(sh|bash|zsh|python|perl|ruby)\b", re.I),
    re.compile(r"xxd\s+-r.*\|.*\b(sh|bash)\b", re.I),
    re.compile(r"printf\s+.*\\x[0-9a-f].*\|\s*(sh|bash)\b", re.I),
    re.compile(r"\$\(.*\bdecode\b.*\)", re.I),
    re.compile(r"curl\s+.*\|\s*(sh|bash|python)\b", re.I),
    re.compile(r"wget\s+.*\|\s*(sh|bash|python)\b", re.I),
]

# Inline eval flags for various interpreters
_INLINE_EVAL_PATTERNS = [
    re.compile(r"\bpython[23]?\s+.*-c\s", re.I),
    re.compile(r"\b(node|bun|deno)\s+.*(-e|--eval)\s", re.I),
    re.compile(r"\bruby\s+.*-e\s", re.I),
    re.compile(r"\bperl\s+.*-[eE]\s", re.I),
    re.compile(r"\bphp\s+.*-r\s", re.I),
]


class PreflightResult:
    """Result of preflight command validation."""

    def __init__(self):
        self.warnings: list[str] = []
        self.blocked: bool = False
        self.block_reason: str | None = None

    def warn(self, msg: str) -> None:
        self.warnings.append(msg)

    def block(self, reason: str) -> None:
        self.blocked = True
        self.block_reason = reason

    @property
    def warning_text(self) -> str:
        return "; ".join(self.warnings) if self.warnings else ""


def preflight_check(command: str, working_dir: str | None = None) -> PreflightResult:
    """Validate command before execution. Detects common model error patterns."""
    result = PreflightResult()

    # Check obfuscation patterns
    for pat in _OBFUSCATION_PATTERNS:
        if pat.search(command):
            result.block(
                "Command blocked: detected encoded/obfuscated execution pattern. "
                "Write commands in plain text instead of encoding them."
            )
            return result

    # Check inline eval (warn, don't block --sometimes legitimate)
    for pat in _INLINE_EVAL_PATTERNS:
        if pat.search(command):
            result.warn("Inline eval detected --prefer writing to a file and executing it")
            break

    # Check for shell variable leakage in script files
    _check_script_variable_leakage(command, working_dir, result)

    return result


def _check_script_variable_leakage(
    command: str, working_dir: str | None, result: PreflightResult,
) -> None:
    """Detect shell $VARIABLE syntax that leaked into Python/JS script files."""
    # Extract script file path from command
    # e.g., "python script.py" → "script.py"
    patterns = [
        re.compile(r"\bpython[23]?\s+([^\s|><;]+\.py)\b"),
        re.compile(r"\b(?:node|bun|deno)\s+([^\s|><;]+\.(?:js|ts|mjs))\b"),
    ]
    for pat in patterns:
        m = pat.search(command)
        if not m:
            continue
        script_path = m.group(1)
        p = Path(script_path)
        if not p.is_absolute() and working_dir:
            p = Path(working_dir) / p
        if not p.is_file() or p.stat().st_size > 512 * 1024:
            continue
        try:
            content = p.read_text(encoding="utf-8", errors="ignore")
        except Exception:
            continue
        shell_vars = _SHELL_VAR_RE.findall(content)
        if len(shell_vars) >= 2:
            examples = ", ".join(shell_vars[:3])
            result.warn(
                f"Shell variable syntax detected in {script_path} ({examples}). "
                "This may indicate shell syntax leaked into the script."
            )


# ── Platform-aware shell resolution ──────────────────────────────────

def resolve_shell() -> tuple[str, list[str]]:
    """Resolve the best shell for the current platform.

    Returns (shell_path, args_prefix).
    """
    system = platform.system()

    if system == "Windows":
        # Prefer PowerShell 7+ (pwsh), fallback to powershell.exe
        pwsh_candidates = [
            r"C:\Program Files\PowerShell\7\pwsh.exe",
            shutil.which("pwsh"),
            shutil.which("powershell"),
        ]
        for candidate in pwsh_candidates:
            if candidate and Path(candidate).is_file():
                return candidate, ["-NoProfile", "-NonInteractive", "-Command"]

        # Last resort: cmd.exe
        cmd = os.environ.get("COMSPEC", "cmd.exe")
        return cmd, ["/c"]

    # POSIX
    shell = os.environ.get("SHELL", "/bin/sh")
    # Fish shell has bashism incompatibilities --fallback
    if "fish" in shell:
        for fallback in ("/bin/bash", "/bin/sh"):
            if Path(fallback).is_file():
                shell = fallback
                break

    return shell, ["-c"]


# ── Process tree cleanup ─────────────────────────────────────────────

async def kill_process_tree(pid: int) -> None:
    """Kill a process and all its children, platform-aware."""
    system = platform.system()

    try:
        if system == "Windows":
            # taskkill /F /T kills the entire process tree
            proc = await asyncio.create_subprocess_exec(
                "taskkill", "/F", "/T", "/PID", str(pid),
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
            )
            await asyncio.wait_for(proc.wait(), timeout=5.0)
        else:
            # Kill process group on Unix
            try:
                os.killpg(os.getpgid(pid), 9)  # SIGKILL
            except (ProcessLookupError, PermissionError):
                try:
                    os.kill(pid, 9)
                except (ProcessLookupError, PermissionError):
                    pass
    except Exception:
        logger.debug("Failed to kill process tree for PID {}", pid, exc_info=True)


# ── Main ExecTool ────────────────────────────────────────────────────

class ExecTool(Toolkit):
    """Shell command execution with accuracy-oriented result enrichment.

    Features:
    - Preflight validation (obfuscation, variable leakage, inline eval)
    - Rich result metadata (exit code classification, duration, cwd, failure reason)
    - Output sanitization (binary filtering, smart truncation)
    - Platform-aware shell selection (PowerShell on Windows, bash/sh on POSIX)
    - Process tree cleanup on timeout
    """

    _MAX_OUTPUT_CHARS = 30_000

    def __init__(
        self,
        timeout: int = 120,
        working_dir: str | None = None,
        deny_patterns: list[str] | None = None,
        allow_patterns: list[str] | None = None,
        restrict_to_workspace: bool = False,
        path_append: str = "",
        max_output_chars: int | None = None,
    ):
        self.timeout = timeout
        self.working_dir = working_dir
        self.deny_patterns = deny_patterns or [
            r"\brm\s+-[rf]{1,2}\b",          # rm -r, rm -rf, rm -fr
            r"\bdel\s+/[fq]\b",              # del /f, del /q
            r"\brmdir\s+/s\b",               # rmdir /s
            r"(?:^|[;&|]\s*)format\b",       # format (as standalone command only)
            r"\b(mkfs|diskpart)\b",          # disk operations
            r"\bdd\s+if=",                   # dd
            r">\s*/dev/sd",                  # write to disk
            r"\b(shutdown|reboot|poweroff)\b",  # system power
            r":\(\)\s*\{.*\};\s*:",          # fork bomb
        ]
        self.allow_patterns = allow_patterns or []
        self.restrict_to_workspace = restrict_to_workspace
        self.path_append = path_append
        self._max_output = max_output_chars or self._MAX_OUTPUT_CHARS
        super().__init__(name="exec", tools=[self.exec])

    async def exec(self, command: str, working_dir: str | None = None, timeout: int | None = None) -> str:
        """Execute a shell command and return enriched results.

        Args:
            command: The shell command to execute.
            working_dir: Working directory (defaults to workspace).
            timeout: Override timeout in seconds for this command.
        """
        cwd = working_dir or self.working_dir or os.getcwd()
        effective_timeout = timeout or self.timeout

        # Step 1: Deny pattern guard
        guard_error = self._guard_command(command, cwd)
        if guard_error:
            return guard_error

        # Step 2: Preflight validation
        preflight = preflight_check(command, cwd)
        if preflight.blocked:
            return f"Error: {preflight.block_reason}"

        # Step 3: Execute
        env = os.environ.copy()
        if self.path_append:
            env["PATH"] = env.get("PATH", "") + os.pathsep + self.path_append
        # Ensure UTF-8 output on Windows
        if platform.system() == "Windows":
            env["PYTHONIOENCODING"] = "utf-8"
            env.setdefault("PYTHONUTF8", "1")

        start_time = time.monotonic()
        timed_out = False

        try:
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=cwd,
                env=env,
            )

            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(),
                    timeout=effective_timeout,
                )
            except asyncio.TimeoutError:
                timed_out = True
                pid = process.pid
                if pid:
                    await kill_process_tree(pid)
                try:
                    await asyncio.wait_for(process.wait(), timeout=5.0)
                except asyncio.TimeoutError:
                    process.kill()
                stdout, stderr = b"", b""

            duration_ms = int((time.monotonic() - start_time) * 1000)
            exit_code = process.returncode

            # Step 4: Build enriched result
            return self._build_result(
                stdout=stdout,
                stderr=stderr,
                exit_code=exit_code,
                duration_ms=duration_ms,
                timed_out=timed_out,
                effective_timeout=effective_timeout,
                cwd=cwd,
                preflight_warnings=preflight.warning_text,
            )

        except Exception as e:
            duration_ms = int((time.monotonic() - start_time) * 1000)
            return self._format_error(str(e), duration_ms, cwd)

    async def execute(self, command: str, working_dir: str | None = None, **kwargs: Any) -> str:
        return await self.exec(command, working_dir, **kwargs)

    # ── Result building ──────────────────────────────────────────────

    def _build_result(
        self,
        stdout: bytes,
        stderr: bytes,
        exit_code: int | None,
        duration_ms: int,
        timed_out: bool,
        effective_timeout: int,
        cwd: str,
        preflight_warnings: str,
    ) -> str:
        """Build enriched result string with metadata."""
        # Decode and sanitize
        stdout_text = sanitize_output(_decode_output(stdout)) if stdout else ""
        stderr_text = sanitize_output(_decode_output(stderr)) if stderr else ""

        # Classify exit code — also check output for command-not-found patterns
        failure_reason = ExitCodeInfo.classify(exit_code, timed_out)
        if failure_reason and exit_code not in (0, None):
            combined_lower = (stdout_text + " " + stderr_text).lower()
            if any(p in combined_lower for p in (
                "not recognized", "not found", "no such file", "command not found",
                "is not recognized as an internal or external command",
                # Chinese Windows: 不是内部或外部命令
                "\u4e0d\u662f\u5185\u90e8\u6216\u5916\u90e8\u547d\u4ee4",
            )):
                failure_reason = "Command not found -- check spelling or install the required tool"
            elif any(p in combined_lower for p in ("permission denied", "access is denied")):
                failure_reason = "Command not executable -- check file permissions"

        # Build output body
        parts: list[str] = []

        if stdout_text.strip():
            parts.append(stdout_text.strip())

        if stderr_text.strip():
            # Only label stderr if there's also stdout
            if parts:
                parts.append(f"STDERR:\n{stderr_text.strip()}")
            else:
                parts.append(stderr_text.strip())

        body = "\n".join(parts) if parts else "(no output)"

        # Truncate if needed
        body, was_truncated = truncate_output(body, self._max_output, strategy="tail")

        # Build metadata header
        meta_parts: list[str] = []

        if failure_reason:
            meta_parts.append(f"[FAILED] {failure_reason}")
            if timed_out:
                meta_parts.append(
                    f"Hint: command timed out after {effective_timeout}s. "
                    f"If expected to take longer, pass timeout={effective_timeout * 3}. "
                    f"For long-running tasks, consider running in background."
                )

        if exit_code is not None and exit_code != 0:
            meta_parts.append(f"Exit code: {exit_code}")

        meta_parts.append(f"Duration: {self._format_duration(duration_ms)}")
        meta_parts.append(f"Working directory: {cwd}")

        if was_truncated:
            meta_parts.append("(output was truncated --use head/tail/grep to inspect specific parts)")

        if preflight_warnings:
            meta_parts.append(f"[Warning] {preflight_warnings}")

        # Combine: metadata at the end so the actual output is first
        meta_block = "\n".join(meta_parts)

        if failure_reason:
            return f"{body}\n\n---\n{meta_block}"
        elif exit_code != 0 or was_truncated or preflight_warnings:
            return f"{body}\n\n---\n{meta_block}"
        else:
            # Success with clean output --minimal metadata
            return body

    def _format_error(self, error: str, duration_ms: int, cwd: str) -> str:
        return (
            f"Error executing command: {error}\n\n---\n"
            f"Duration: {self._format_duration(duration_ms)}\n"
            f"Working directory: {cwd}"
        )

    @staticmethod
    def _format_duration(ms: int) -> str:
        if ms < 1000:
            return f"{ms}ms"
        elif ms < 60_000:
            return f"{ms / 1000:.1f}s"
        else:
            minutes = ms // 60_000
            secs = (ms % 60_000) / 1000
            return f"{minutes}m{secs:.0f}s"

    # ── Command guards ───────────────────────────────────────────────

    def _guard_command(self, command: str, cwd: str) -> str | None:
        """Best-effort safety guard for potentially destructive commands."""
        cmd = command.strip()
        lower = cmd.lower()

        for pattern in self.deny_patterns:
            if re.search(pattern, lower):
                return "Error: Command blocked by safety guard (dangerous pattern detected)"

        if self.allow_patterns:
            if not any(re.search(p, lower) for p in self.allow_patterns):
                return "Error: Command blocked by safety guard (not in allowlist)"

        if self.restrict_to_workspace:
            if "..\\" in cmd or "../" in cmd:
                return "Error: Command blocked by safety guard (path traversal detected)"

            cwd_path = Path(cwd).resolve()

            for raw in self._extract_absolute_paths(cmd):
                try:
                    p = Path(raw.strip()).resolve()
                except Exception:
                    continue
                if p.is_absolute() and cwd_path not in p.parents and p != cwd_path:
                    return "Error: Command blocked by safety guard (path outside working dir)"

        return None

    @staticmethod
    def _extract_absolute_paths(command: str) -> list[str]:
        win_paths = re.findall(r"[A-Za-z]:\\[^\s\"'|><;]+", command)   # Windows: C:\...
        posix_paths = re.findall(r"(?:^|[\s|>])(/[^\s\"'>]+)", command) # POSIX: /absolute only
        return win_paths + posix_paths
