"""Simplified middleware that exposes a basic shell tool to agents."""

from __future__ import annotations

import logging
import os
import re
import subprocess
import sys
from typing import Any

from langchain.agents.middleware.types import AgentMiddleware, AgentState
from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from langchain_core.tools.base import ToolException

# Detect if running on Windows
IS_WINDOWS = sys.platform == "win32"

# Security: Audit logger for shell commands
_shell_audit_logger = logging.getLogger("ieum.security.shell")

# Security: Blocked commands that are obviously destructive and never part of
# normal development workflows. This is intentionally a short list -- the agent
# needs broad shell access to function as a coding assistant.
BLOCKED_COMMANDS: list[str] = [
    "rm -rf /",
    "rm -rf /*",
    "rm -rf --no-preserve-root /",
    "mkfs",
    "mkfs.",  # mkfs.ext4, mkfs.xfs, etc.
    "dd if=/dev/zero",
    "dd if=/dev/random",
    "dd if=/dev/urandom",
    ":(){ :|:& };:",  # fork bomb
    "> /dev/sda",
    "chmod -R 777 /",
    "chown -R",
]

# Regex patterns for dangerous command constructs
BLOCKED_PATTERNS: list[re.Pattern[str]] = [
    # Fork bomb variants
    re.compile(r":\(\)\s*\{.*\|.*&\s*\}\s*;?\s*:"),
    # Pipe to shell for remote code execution
    re.compile(r"curl\s.*\|\s*(ba)?sh"),
    re.compile(r"wget\s.*\|\s*(ba)?sh"),
    re.compile(r"curl\s.*\|\s*sudo\s+(ba)?sh"),
    re.compile(r"wget\s.*\|\s*sudo\s+(ba)?sh"),
    # Overwrite entire disk
    re.compile(r">\s*/dev/[sh]d[a-z]"),
    re.compile(r"dd\s+if=.*\s+of=/dev/[sh]d[a-z]"),
    # Recursive remove from root
    re.compile(r"rm\s+(-[a-zA-Z]*f[a-zA-Z]*\s+)*(-[a-zA-Z]*r[a-zA-Z]*\s+)*\s*/\s*$"),
    re.compile(r"rm\s+(-[a-zA-Z]*r[a-zA-Z]*\s+)*(-[a-zA-Z]*f[a-zA-Z]*\s+)*\s*/\s*$"),
]


def _is_command_blocked(command: str, *, safe_mode: bool = False) -> str | None:
    """Check if a command should be blocked.

    Args:
        command: The shell command to validate.
        safe_mode: If True, apply stricter validation rules.

    Returns:
        A reason string if blocked, or None if the command is allowed.
    """
    cmd_stripped = command.strip()
    cmd_lower = cmd_stripped.lower()

    # Check against blocked command literals
    for blocked in BLOCKED_COMMANDS:
        if blocked.lower() in cmd_lower:
            return f"Blocked dangerous command pattern: {blocked}"

    # Check against blocked regex patterns
    for pattern in BLOCKED_PATTERNS:
        if pattern.search(cmd_stripped):
            return f"Blocked dangerous command pattern matching: {pattern.pattern}"

    # Safe mode: additional restrictions
    if safe_mode:
        # Block sudo usage entirely in safe mode
        if re.search(r"\bsudo\b", cmd_stripped):
            return "sudo is not allowed in safe mode"
        # Block modifying system directories
        if re.search(r"\b(rm|mv|cp)\b.*\s+/(usr|etc|var|boot|sys|proc)\b", cmd_stripped):
            return "Modifying system directories is not allowed in safe mode"

    return None


class ShellMiddleware(AgentMiddleware[AgentState, Any]):
    """Give basic shell access to agents via the shell.

    This shell will execute on the local machine and has NO safeguards except
    for the human in the loop safeguard provided by the CLI itself.

    On Windows, commands are executed via PowerShell for better compatibility.
    On Unix-like systems, commands are executed via the default shell.
    """

    def __init__(
        self,
        *,
        workspace_root: str,
        timeout: float = 120.0,
        max_output_bytes: int = 100_000,
        env: dict[str, str] | None = None,
        safe_mode: bool = False,
    ) -> None:
        """Initialize an instance of `ShellMiddleware`.

        Args:
            workspace_root: Working directory for shell commands.
            timeout: Maximum time in seconds to wait for command completion.
                Defaults to 120 seconds.
            max_output_bytes: Maximum number of bytes to capture from command output.
                Defaults to 100,000 bytes.
            env: Environment variables to pass to the subprocess. If None,
                uses the current process's environment. Defaults to None.
            safe_mode: If True, enable stricter command validation that blocks
                sudo and system directory modifications. Defaults to False.
        """
        super().__init__()
        self._timeout = timeout
        self._max_output_bytes = max_output_bytes
        self._tool_name = "shell"
        self._env = env if env is not None else os.environ.copy()
        self._workspace_root = workspace_root
        self._safe_mode = safe_mode or os.environ.get("IEUM_SAFE_MODE", "").lower() == "true"

        # Build description with working directory and platform information
        if IS_WINDOWS:
            shell_info = "PowerShell"
            description = (
                f"Execute a PowerShell command directly on the host. Commands will run in "
                f"the working directory: {workspace_root}. Each command runs in a fresh PowerShell "
                f"environment. Use PowerShell syntax and cmdlets (Get-ChildItem, Get-Content, etc.) "
                f"or standard Windows commands. Commands may be truncated if they exceed the "
                f"configured timeout or output limits."
            )
        else:
            shell_info = "bash/sh"
            description = (
                f"Execute a shell command directly on the host. Commands will run in "
                f"the working directory: {workspace_root}. Each command runs in a fresh shell "
                f"environment with the current process's environment variables. Commands may "
                f"be truncated if they exceed the configured timeout or output limits."
            )

        @tool(self._tool_name, description=description)
        def shell_tool(
            command: str,
            runtime: ToolRuntime[None, AgentState],
        ) -> ToolMessage | str:
            """Execute a shell command.

            Args:
                command: The shell command to execute.
                runtime: The tool runtime context.
            """
            return self._run_shell_command(command, tool_call_id=runtime.tool_call_id)

        self._shell_tool = shell_tool
        self.tools = [self._shell_tool]

    def _run_shell_command(
        self,
        command: str,
        *,
        tool_call_id: str | None,
    ) -> ToolMessage | str:
        """Execute a shell command and return the result.

        Args:
            command: The shell command to execute.
            tool_call_id: The tool call ID for creating a ToolMessage.

        Returns:
            A ToolMessage with the command output or an error message.

        Note:
            On Windows, commands are executed via PowerShell for better compatibility.
            On Unix-like systems, commands are executed via the default shell (shell=True).
        """
        if not command or not isinstance(command, str):
            msg = "Shell tool expects a non-empty command string."
            raise ToolException(msg)

        # Security: Audit log every command execution
        _shell_audit_logger.info(
            "Shell command executed | cwd=%s | command=%s",
            self._workspace_root,
            command,
        )

        # Security: Validate command against blocked patterns
        safe_mode = self._safe_mode
        block_reason = _is_command_blocked(command, safe_mode=safe_mode)
        if block_reason is not None:
            _shell_audit_logger.warning(
                "Shell command BLOCKED | reason=%s | command=%s",
                block_reason,
                command,
            )
            return ToolMessage(
                content=f"Error: Command blocked by security policy. {block_reason}",
                tool_call_id=tool_call_id,
                name=self._tool_name,
                status="error",
            )

        try:
            if IS_WINDOWS:
                # Use PowerShell on Windows for better compatibility
                # -NoProfile: Skip loading user profile for faster startup
                # -ExecutionPolicy Bypass: Allow script execution
                # -Command: Execute the provided command
                result = subprocess.run(
                    [
                        "powershell",
                        "-NoProfile",
                        "-ExecutionPolicy", "Bypass",
                        "-Command",
                        command,
                    ],
                    check=False,
                    capture_output=True,
                    text=True,
                    timeout=self._timeout,
                    env=self._env,
                    cwd=self._workspace_root,
                )
            else:
                # Use default shell on Unix-like systems
                result = subprocess.run(
                    command,
                    check=False,
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=self._timeout,
                    env=self._env,
                    cwd=self._workspace_root,
                )

            # Combine stdout and stderr
            output_parts = []
            if result.stdout:
                output_parts.append(result.stdout)
            if result.stderr:
                stderr_lines = result.stderr.strip().split("\n")
                for line in stderr_lines:
                    output_parts.append(f"[stderr] {line}")

            output = "\n".join(output_parts) if output_parts else "<no output>"

            # Truncate output if needed
            if len(output) > self._max_output_bytes:
                output = output[: self._max_output_bytes]
                output += f"\n\n... Output truncated at {self._max_output_bytes} bytes."

            # Add exit code info if non-zero
            if result.returncode != 0:
                output = f"{output.rstrip()}\n\nExit code: {result.returncode}"
                status = "error"
            else:
                status = "success"

        except subprocess.TimeoutExpired:
            output = f"Error: Command timed out after {self._timeout:.1f} seconds."
            status = "error"

        return ToolMessage(
            content=output,
            tool_call_id=tool_call_id,
            name=self._tool_name,
            status=status,
        )


__all__ = ["ShellMiddleware"]
