# -*- coding: utf-8 -*-
# flake8: noqa: E501
# pylint: disable=line-too-long
"""The shell command tool."""

import asyncio
import locale
import os
import re
import signal
import subprocess
import sys
import tempfile
from dataclasses import replace
from pathlib import Path
from typing import Any, Optional

from agentscope.message import TextBlock, ToolResultState
from agentscope.tool import ToolChunk

from ...config.context import (
    get_current_shell_command_executable,
    get_current_shell_command_timeout,
    get_current_workspace_dir,
)
from ...constant import WORKING_DIR
from ...runtime.tool_registry import tool_descriptor
from ...sandbox import ExecutionResult


def _kill_process_tree_win32(pid: int) -> None:
    """Kill a process and all its descendants on Windows via taskkill.

    Uses ``taskkill /F /T`` which forcefully terminates the entire process
    tree, including grandchild processes that ``Popen.kill()`` would miss.
    """
    try:
        subprocess.call(
            ["taskkill", "/F", "/T", "/PID", str(pid)],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            timeout=10,
        )
    except Exception:
        pass


def _windows_shell_creationflags() -> int:
    """Return Windows process flags for shell commands."""
    return getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)


def _collapse_newlines_outside_quotes(cmd: str) -> str:
    r"""Collapse newlines outside quoted strings; preserve those inside.

    Used only on Unix where sh/bash correctly handles newlines in quotes.
    Handles backslash-newline (line continuation) by removing both chars,
    and treats single-quoted content as fully literal per POSIX.
    """
    result: list[str] = []
    in_single_quote = False
    in_double_quote = False
    i = 0
    length = len(cmd)

    while i < length:
        char = cmd[i]

        # Toggle quote state
        if char == "'" and not in_double_quote:
            in_single_quote = not in_single_quote
            result.append(char)
            i += 1
            continue

        if char == '"' and not in_single_quote:
            in_double_quote = not in_double_quote
            result.append(char)
            i += 1
            continue

        # Inside single quotes: everything is literal (POSIX)
        if in_single_quote:
            result.append(char)
            i += 1
            continue

        # Backslash-newline (line continuation): remove both chars
        if char == "\\" and i + 1 < length and cmd[i + 1] in ("\r", "\n"):
            i += 2
            # \r\n sequence: skip the \n as well
            if i < length and cmd[i - 1] == "\r" and cmd[i] == "\n":
                i += 1
            continue

        # Backslash escape (non-newline): keep both chars
        if char == "\\" and i + 1 < length:
            result.append(char)
            result.append(cmd[i + 1])
            i += 2
            continue

        # Newlines
        if char in ("\r", "\n"):
            if in_double_quote:
                # Preserve newlines inside double quotes
                result.append(char)
            else:
                # Collapse \r\n as a single space
                if char == "\r" and i + 1 < length and cmd[i + 1] == "\n":
                    i += 1
                result.append(" ")
            i += 1
            continue

        result.append(char)
        i += 1

    return "".join(result)


def _collapse_embedded_newlines(cmd: str) -> str:
    r"""Replace embedded newline characters with spaces in a command string.

    LLMs produce tool-call arguments in JSON where ``\n`` is parsed as an
    actual newline character.  In the original shell command the user
    intended the *literal* two-character sequence ``\n`` (e.g. inside a
    ``--content`` flag), but after JSON decoding it becomes a real line
    break.  When passed to a shell:

    * **Windows** ``cmd.exe`` truncates the command at the first newline
      regardless of quoting context — this is a hard limitation of the
      Windows command processor.  All newlines must be collapsed.
    * **Unix** ``sh -c`` treats an unquoted newline as a command separator,
      but correctly handles newlines inside quoted strings.

    On Unix/macOS, newlines inside quoted strings are preserved so that
    downstream commands receive the correct multi-line content (e.g.
    ``--text "Hello\nWorld"``).  On Windows, all newlines are collapsed
    to ensure the command at least executes successfully.
    """
    if "\n" not in cmd:
        return cmd
    if sys.platform == "win32":
        # cmd.exe truncates at newlines regardless of quoting — must
        # collapse all to ensure the command executes at all.
        return cmd.replace("\r\n", " ").replace("\n", " ")
    return _collapse_newlines_outside_quotes(cmd)


def _sanitize_win_cmd(cmd: str) -> str:
    """Fix common LLM escaping artefacts for Windows ``cmd.exe``.

    LLMs sometimes produce commands with backslash-escaped double quotes
    (``\\"``) — valid in bash/JSON but meaningless to ``cmd.exe``.  When
    *every* double-quote in the command is preceded by a backslash, it is
    almost certainly a double-escape artefact, so we strip them.
    """
    if '\\"' in cmd and '"' not in cmd.replace('\\"', ""):
        return cmd.replace('\\"', '"')
    return cmd


def _read_temp_file(path: str) -> str:
    """Read a temporary output file and return its decoded content."""
    try:
        with open(path, "rb") as f:
            return smart_decode(f.read())
    except OSError:
        return ""


def _shell_basename(executable: str) -> str:
    """Extract lowercase basename from a path using both / and \\ separators."""
    return executable.replace("\\", "/").rsplit("/", 1)[-1].lower()


def _is_powershell(executable: str) -> bool:
    """Check if the given executable path is a PowerShell variant."""
    return _shell_basename(executable) in (
        "powershell",
        "powershell.exe",
        "pwsh",
        "pwsh.exe",
    )


def _is_cmd(executable: str) -> bool:
    """Check if the given executable path is cmd.exe."""
    return _shell_basename(executable) in ("cmd", "cmd.exe")


_PS_CMD_RE = re.compile(
    r"^(powershell(?:\.exe)?|pwsh(?:\.exe)?)"
    r"((?:\s+-(?:NoProfile|NonInteractive|NoLogo))*)"
    r"(?:\s+-ExecutionPolicy\s+\S+)?"
    r"\s+-Command\s+",
    re.IGNORECASE,
)


def _extract_powershell_command(cmd: str) -> tuple[str | None, str]:
    """Detect ``powershell -Command <body>`` and return (exe, inner_body).

    When *cmd* starts with a PowerShell invocation followed by ``-Command``,
    extract the executable name and the inner command body (with a single
    layer of surrounding double-quotes removed if present).

    Returns ``(None, cmd)`` unchanged when no PowerShell prefix is found.
    """
    m = _PS_CMD_RE.match(cmd)
    if not m:
        return None, cmd
    ps_exe = m.group(1)
    inner = cmd[m.end() :]
    if len(inner) >= 2 and inner[0] == '"' and inner[-1] == '"':
        inner = inner[1:-1]
    return ps_exe, inner


# pylint: disable=too-many-branches, too-many-statements
def _execute_subprocess_sync(
    cmd: str,
    cwd: str,
    timeout: float,
    env: dict | None = None,
    shell_executable: str | None = None,
) -> tuple[int, str, str]:
    """Execute subprocess synchronously in a thread.

    This function runs in a separate thread to avoid Windows asyncio
    subprocess limitations.

    stdout/stderr are redirected to temporary files instead of pipes.
    On Windows, child processes inherit pipe handles and keep them open
    even after the parent exits, which causes ``communicate()`` to block
    until *all* holders close (e.g. a Chrome process launched via
    ``Start-Process``).  With temp-file redirection, ``proc.wait()``
    only waits for the direct child (``cmd.exe``) to exit, so commands
    that spawn background processes return immediately.

    .. note::

       Callers must pre-process *cmd* through
       :func:`_collapse_embedded_newlines` before passing it here.
       ``execute_shell_command`` already does this.

    Args:
        cmd (`str`):
            The shell command to execute (must not contain embedded
            newlines — see note above).
        cwd (`str`):
            The working directory for the command execution.
        timeout (`float`):
            The maximum time (in seconds) allowed for the command to run.
        env (`dict | None`):
            Environment variables for the subprocess.
        shell_executable (`str | None`):
            Path to the shell executable. When ``None``, defaults to
            ``cmd.exe``.

    Returns:
        `tuple[int, str, str]`:
            A tuple containing the return code, standard output, and
            standard error of the executed command. If timeout occurs, the
            return code will be -1 and stderr will contain timeout information.
    """
    stdout_path: str | None = None
    stderr_path: str | None = None
    stdout_file = None
    stderr_file = None

    try:
        if shell_executable and _is_powershell(shell_executable):
            # Strip redundant powershell/pwsh -Command wrapper that the
            # LLM may emit even though the shell is already PowerShell.
            _, cmd = _extract_powershell_command(cmd)
            wrapped = [
                shell_executable,
                "-NoProfile",
                "-NonInteractive",
                "-Command",
                cmd,
            ]
        elif not shell_executable or _is_cmd(shell_executable):
            cmd = _sanitize_win_cmd(cmd)
            shell_name = shell_executable or "cmd"
            wrapped = f'{shell_name} /D /S /C "{cmd}"'
        else:
            # POSIX-like shell on Windows (e.g. Git Bash, MSYS2)
            wrapped = [shell_executable, "-c", cmd]

        stdout_fd, stdout_path = tempfile.mkstemp(prefix="qwenpaw_out_")
        stderr_fd, stderr_path = tempfile.mkstemp(prefix="qwenpaw_err_")
        stdout_file = os.fdopen(stdout_fd, "wb")
        stderr_file = os.fdopen(stderr_fd, "wb")

        proc = subprocess.Popen(  # pylint: disable=consider-using-with
            wrapped,
            shell=False,
            stdout=stdout_file,
            stderr=stderr_file,
            text=False,
            cwd=cwd,
            env=env,
            creationflags=_windows_shell_creationflags(),
        )

        # Parent copies are no longer needed — the child inherited its own
        # handles via CreateProcess.  Closing here avoids holding the files
        # open longer than necessary.
        stdout_file.close()
        stdout_file = None
        stderr_file.close()
        stderr_file = None

        timed_out = False
        try:
            proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            timed_out = True
            _kill_process_tree_win32(proc.pid)
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                try:
                    proc.kill()
                except OSError:
                    pass

        stdout_str = _read_temp_file(stdout_path)
        stderr_str = _read_temp_file(stderr_path)

        if timed_out:
            timeout_msg = (
                f"Command execution exceeded the timeout of {timeout} seconds."
            )
            if stderr_str:
                stderr_str = f"{stderr_str}\n{timeout_msg}"
            else:
                stderr_str = timeout_msg
            return -1, stdout_str, stderr_str

        returncode = proc.returncode if proc.returncode is not None else -1
        return returncode, stdout_str, stderr_str

    except Exception as e:
        return -1, "", str(e)
    finally:
        for f in (stdout_file, stderr_file):
            if f is not None:
                try:
                    f.close()
                except OSError:
                    pass
        for path in (stdout_path, stderr_path):
            if path is not None:
                try:
                    os.unlink(path)
                except OSError:
                    pass


# Extra seconds added to the tool-call deadline to accommodate first-time
# sandbox creation (user provisioning, profile creation, firewall rules, ACLs).
# Subsequent calls hit the cache and need no extension.
_SANDBOX_SETUP_DEADLINE_EXTENSION = 180.0


async def _execute_in_sandbox(
    cmd: str,
    sandbox_config: Any,
    timeout: float,
    cwd: str,
    env: dict[str, str],
) -> ExecutionResult:
    """Execute a shell command inside the sandbox and return raw result.

    On first invocation the sandbox setup (user creation, profile, ACLs,
    firewall rules) can take 5-100+ seconds. To prevent the ToolCoordinator's
    deadline from expiring during this one-time setup, we temporarily extend
    the deadline by _SANDBOX_SETUP_DEADLINE_EXTENSION seconds. The extension
    is only applied when the call context has a deadline set.
    """
    from ...sandbox import create_sandbox
    from ...tool_calls import get_call_context

    # Sandbox backends rebuild their environment from os.environ. Carry over
    # the PATH adjusted by the shell entrypoint unless policy set one itself.
    sandbox_env = dict(sandbox_config.env_vars)
    if not any(key.upper() == "PATH" for key in sandbox_env):
        path_key = next(
            (key for key in env if key.upper() == "PATH"),
            "PATH",
        )
        sandbox_env[path_key] = env[path_key]

    effective_config = replace(
        sandbox_config,
        timeout_seconds=int(timeout),
        env_vars=sandbox_env,
    )

    # Temporarily extend the tool-call deadline so that sandbox creation
    # does not consume the user's command timeout budget.
    ctx = get_call_context()
    original_deadline = None
    if ctx is not None and ctx.deadline is not None:
        original_deadline = ctx.deadline
        ctx.deadline += _SANDBOX_SETUP_DEADLINE_EXTENSION

    try:
        async with create_sandbox(effective_config) as sandbox:
            # Restore the original deadline (plus only the command timeout)
            # now that sandbox setup is complete.
            if ctx is not None and original_deadline is not None:
                now = asyncio.get_event_loop().time()
                ctx.deadline = now + timeout
            result = await sandbox.execute(cmd, cwd=cwd)
    except BaseException:
        # On failure, restore original deadline to avoid permanent extension
        if ctx is not None and original_deadline is not None:
            ctx.deadline = original_deadline
        raise

    return result


_DANGER_NAMES = {
    "python",
    "pythonw",
    "cmd",
    "powershell",
    "pwsh",
    "conhost",
}

# Prefix: kill/taskkill at command start or after &&, ;, |
_KILL_PREFIX = r"(?:^|[;&|]\s*)\s*"

# Matches PID-based kills: taskkill /PID 123, kill -9 123, kill 123.
_KILL_PID_RE = re.compile(
    rf"{_KILL_PREFIX}(?:taskkill|kill|stop-process)\b"
    rf".*(?:/PID|-p|-pid|\b)\s*(\d+)",
    re.IGNORECASE,
)

# Matches dangerous process names as /IM targets or bare kill targets.
_DANGER_NAME_RE = re.compile(
    rf"{_KILL_PREFIX}(?:taskkill|kill|stop-process)\b"
    rf".*?\b({'|'.join(_DANGER_NAMES)})(?:\.exe)?\b",
    re.IGNORECASE,
)

# Shell variables that reference the current/parent PID.
_SHELL_PID_VARS = {"$$", "$ppid", "$pid"}


def _is_dangerous_self_kill(cmd: str) -> bool:
    """Return True if *cmd* would kill the current process or its parent.

    Uses token-based regex matching to avoid false positives from
    substring matching (e.g. ``echo "do not kill python"`` is safe).

    Blocks three patterns:
    1. ``taskkill /IM <dangerous_name>`` — kills by image name.
    2. ``kill <pid>`` / ``taskkill /PID <pid>`` targeting our PID or
       parent.
    3. Shell variable self-kill: ``kill -9 $$``, ``kill $PPID``.
    """
    lower = cmd.lower()

    if _DANGER_NAME_RE.search(lower):
        return True

    if "kill" in lower or "stop-process" in lower:
        if any(var in lower for var in _SHELL_PID_VARS):
            return True

    m = _KILL_PID_RE.search(lower)
    if m:
        try:
            target_pid = int(m.group(1))
            protected_pids = {os.getpid()}
            if hasattr(os, "getppid"):
                protected_pids.add(os.getppid())
            if target_pid in protected_pids:
                return True
        except ValueError:
            pass

    return False


# pylint: disable=too-many-branches, too-many-statements
@tool_descriptor(
    requires_sandbox=("shell_exec",),
    async_execution=True,
    tool_type="shell",
    target_param="command",
    policy_name="Bash",
    ui_description="Execute shell commands",
    ui_icon="💻",
)
async def execute_shell_command(
    command: str,
    timeout: float = 60.0,
    cwd: Optional[Path] = None,
    sandbox_config: Optional[Any] = None,
) -> ToolChunk:
    """Execute a shell command and return its output.

    Each call runs in a fresh subprocess — `cd`, `export`, `source`,
    etc. do NOT persist. Pass `cwd=` or chain in one call
    (`cd /repo && pytest`).

    IMPORTANT: Check the 'Default Shell' field to
    determine which shell is active, and generate commands using the
    appropriate syntax (e.g. bash vs PowerShell vs cmd.exe).

    Args:
        command (`str`):
            The shell command to execute.
        timeout (`float`, defaults to `60.0`):
            The maximum time (in seconds) allowed for the command to run.
            Default is 60.0 seconds.
        cwd (`Optional[Path]`, defaults to `None`):
            The working directory for the command execution.
            If None, defaults to the agent workspace.
        sandbox_config (`Optional[Any]`, defaults to `None`):
            Sandbox execution configuration compiled from governance policy.
            When provided, the command executes within a sandboxed environment
            with the specified mount permissions and network restrictions.

    Returns:
        `ToolChunk`:
            The tool response containing the return code, standard output, and
            standard error of the executed command. If timeout occurs, the
            return code will be -1 and stderr will contain timeout information.
    """

    cmd = _collapse_embedded_newlines((command or "").strip())

    if _is_dangerous_self_kill(cmd):
        return ToolChunk(
            is_last=True,
            state=ToolResultState.ERROR,
            content=[
                TextBlock(
                    type="text",
                    text=(
                        "Blocked: this command would terminate the "
                        "QwenPaw process or its parent. "
                        "Refusing to execute."
                    ),
                ),
            ],
        )

    if isinstance(timeout, str):
        try:
            timeout = float(timeout)
        except (ValueError, TypeError):
            timeout = 60.0

    # Apply agent-configured default when the caller used the hardcoded
    # default (60.0).  An explicit LLM-provided value != 60.0 is kept.
    if timeout == 60.0:
        configured = get_current_shell_command_timeout()
        if configured is not None:
            timeout = configured

    # Use current workspace_dir from context, fallback to WORKING_DIR
    if cwd is not None:
        working_dir = cwd
    else:
        working_dir = get_current_workspace_dir() or WORKING_DIR

    # Ensure the venv Python is on PATH for subprocesses
    env = os.environ.copy()
    python_bin_dir = str(Path(sys.executable).parent)
    existing_path = env.get("PATH", "")
    if existing_path:
        env["PATH"] = python_bin_dir + os.pathsep + existing_path
    else:
        env["PATH"] = python_bin_dir

    shell_executable = (
        get_current_shell_command_executable()
        or os.environ.get("SHELL")
        or None
    )

    if sandbox_config is not None:
        # Create a copy with resolved shell and timeout to avoid mutating
        # the shared config object (it may be reused across tool calls).
        sandbox_config = replace(
            sandbox_config,
            shell_executable=shell_executable,
            timeout_seconds=int(timeout),
        )
        result = await _execute_in_sandbox(
            cmd,
            sandbox_config,
            timeout,
            str(working_dir),
            env,
        )
        # Sandbox violation: command tried to access something not permitted
        if result.sandbox_violation:
            return ToolChunk(
                is_last=True,
                state=ToolResultState.DENIED,
                content=[
                    TextBlock(
                        type="text",
                        text=f"Sandbox violation: {result.sandbox_violation}\n"
                        f"Command was blocked by sandbox security policy.",
                    ),
                ],
                metadata={"sandbox_violation": result.sandbox_violation},
            )
        if result.exit_code == 0:
            response_text = (
                result.stdout or "Command executed successfully (no output)."
            )
            if result.stderr:
                response_text += f"\n[stderr]\n{result.stderr}"
        else:
            parts = [f"Command failed with exit code {result.exit_code}."]
            if result.stdout:
                parts.append(f"\n[stdout]\n{result.stdout}")
            if result.stderr:
                parts.append(f"\n[stderr]\n{result.stderr}")
            response_text = "".join(parts)
        return ToolChunk(
            is_last=True,
            state=ToolResultState.SUCCESS,
            content=[
                TextBlock(
                    type="text",
                    text=response_text,
                ),
            ],
        )

    import logging as _logging

    _logging.getLogger(__name__).debug(
        "[sandbox] SKIP: sandbox_config is None, executing directly",
    )

    try:
        if sys.platform == "win32":
            # Windows: use thread pool to avoid asyncio subprocess limitations
            returncode, stdout_str, stderr_str = await asyncio.to_thread(
                _execute_subprocess_sync,
                cmd,
                str(working_dir),
                timeout,
                env,
                shell_executable,
            )
        else:
            proc = await asyncio.create_subprocess_shell(
                cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                bufsize=0,
                cwd=str(working_dir),
                env=env,
                start_new_session=True,
                executable=shell_executable,
            )

            try:
                # Apply timeout to communicate directly; wait()+communicate()
                # can hang if descendants keep stdout/stderr pipes open.
                from ...tool_calls import cancellable_wait

                stdout, stderr = await cancellable_wait(
                    proc.communicate(),
                    fallback_secs=timeout,
                )
                stdout_str = smart_decode(stdout)
                stderr_str = smart_decode(stderr)
                returncode = proc.returncode

            except (asyncio.TimeoutError, asyncio.CancelledError):
                stderr_suffix = (
                    f"⚠️ TimeoutError: The command execution exceeded "
                    f"the timeout of {timeout} seconds. "
                    f"Please consider increasing the timeout value if this command "
                    f"requires more time to complete."
                )
                returncode = -1
                try:
                    # Kill the entire process group so that child processes
                    # spawned by the shell are also terminated.
                    pgid = os.getpgid(proc.pid)
                    os.killpg(pgid, signal.SIGTERM)
                    try:
                        await asyncio.wait_for(proc.wait(), timeout=2)
                    except asyncio.TimeoutError:
                        os.killpg(pgid, signal.SIGKILL)
                        await asyncio.wait_for(proc.wait(), timeout=2)

                    # Drain remaining output.
                    try:
                        stdout, stderr = await asyncio.wait_for(
                            proc.communicate(),
                            timeout=1,
                        )
                    except asyncio.TimeoutError:
                        stdout, stderr = b"", b""
                    stdout_str = smart_decode(stdout)
                    stderr_str = smart_decode(stderr)
                    if stderr_str:
                        stderr_str += f"\n{stderr_suffix}"
                    else:
                        stderr_str = stderr_suffix
                except (ProcessLookupError, OSError):
                    # Process already gone or pgid lookup failed — fall back
                    # to direct kill on the process itself.
                    try:
                        proc.kill()
                        await proc.wait()
                    except (ProcessLookupError, OSError):
                        pass
                    stdout_str = ""
                    stderr_str = stderr_suffix

        if returncode == 0:
            if stdout_str:
                response_text = stdout_str
            else:
                response_text = "Command executed successfully (no output)."
            if stderr_str:
                response_text += f"\n[stderr]\n{stderr_str}"
        else:
            response_parts = [f"Command failed with exit code {returncode}."]
            if stdout_str:
                response_parts.append(f"\n[stdout]\n{stdout_str}")
            if stderr_str:
                response_parts.append(f"\n[stderr]\n{stderr_str}")
            response_text = "".join(response_parts)

        return ToolChunk(
            is_last=True,
            state=ToolResultState.SUCCESS,
            content=[
                TextBlock(
                    type="text",
                    text=response_text,
                ),
            ],
        )

    except Exception as e:
        return ToolChunk(
            is_last=True,
            state=ToolResultState.SUCCESS,
            content=[
                TextBlock(
                    type="text",
                    text=f"Error: Shell command execution failed due to \n{e}",
                ),
            ],
        )


def smart_decode(data: bytes) -> str:
    try:
        decoded_str = data.decode("utf-8")
    except UnicodeDecodeError:
        encoding = locale.getpreferredencoding(False) or "utf-8"
        decoded_str = data.decode(encoding, errors="replace")

    return decoded_str.strip("\n")
