"""Useful tools for agents: file editing and bash execution."""

import logging
import os
import re
import shlex
import shutil
import signal
import subprocess
import sys
import threading
from collections.abc import Callable
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


def _log_exc() -> None:
    logger.debug("Exception caught", exc_info=True)


def _find_windows_bash() -> str | None:
    """Find bash.exe on Windows (Git for Windows, WSL, etc.)."""
    found = shutil.which("bash")
    if found:
        return found
    for candidate in [
        r"C:\Program Files\Git\bin\bash.exe",
        r"C:\Program Files\Git\usr\bin\bash.exe",
        r"C:\Program Files (x86)\Git\bin\bash.exe",
    ]:
        if Path(candidate).exists():
            return candidate
    return None


_WINDOWS_BASH: str | None = _find_windows_bash() if sys.platform == "win32" else None


def _popen_kwargs(command: str) -> dict[str, Any]:
    """Return Popen kwargs appropriate for the current platform.

    On Unix, uses ``shell=True`` with ``start_new_session=True``.
    On Windows with bash available, invokes bash directly.
    On Windows without bash, falls back to PowerShell.

    Args:
        command: The command string to execute.

    Returns:
        Dict of keyword arguments for ``subprocess.Popen``.
    """
    if sys.platform != "win32":
        return {
            "args": command,
            "shell": True,
            "start_new_session": True,
        }
    if _WINDOWS_BASH:
        return {
            "args": [_WINDOWS_BASH, "-c", command],
            "creationflags": subprocess.CREATE_NEW_PROCESS_GROUP,
        }
    ps = shutil.which("pwsh") or shutil.which("powershell") or "powershell"
    return {
        "args": [ps, "-NoProfile", "-Command", command],
        "creationflags": subprocess.CREATE_NEW_PROCESS_GROUP,
    }


def _truncate_output(output: str, max_chars: int) -> str:
    if len(output) <= max_chars:
        return output
    worst_msg = f"\n\n... [truncated {len(output)} chars] ...\n\n"
    if max_chars < len(worst_msg):
        return output[:max_chars]
    remaining = max_chars - len(worst_msg)
    head = remaining // 2
    tail = remaining - head
    dropped = len(output) - head - tail
    msg = f"\n\n... [truncated {dropped} chars] ...\n\n"
    if tail:
        return output[:head] + msg + output[-tail:]
    return output[:head] + msg


DISALLOWED_BASH_COMMANDS = {
    ".",
    "env",
    "eval",
    "exec",
    "source",
}

DISALLOWED_PS_COMMANDS = {
    "Invoke-Expression",
    "iex",
}


_SHELL_PREFIX_TOKENS = frozenset(("!", "{", "}", "(", ")", "&"))
_REDIRECT_RE = re.compile(r"^[0-9]*[<>][<>&]*")


def _extract_leading_command_name(part: str) -> str | None:
    try:
        tokens = shlex.split(part)
    except ValueError:
        _log_exc()
        return None
    if not tokens:
        return None

    i = 0
    while i < len(tokens) and re.match(r"^[A-Za-z_][A-Za-z0-9_]*=.*", tokens[i]):
        i += 1

    while i < len(tokens):
        token = tokens[i]
        if token in _SHELL_PREFIX_TOKENS:
            i += 1
            continue
        m = _REDIRECT_RE.match(token)
        if m:
            if m.end() < len(token):
                i += 1
            else:
                i += 2
            continue
        break

    if i >= len(tokens):
        return None
    name = tokens[i].lstrip("({")
    if not name:
        return None
    return name.split("/")[-1]


def _split_respecting_quotes(command: str, pattern: re.Pattern[str]) -> list[str]:
    """Split *command* on *pattern* while skipping quoted and escaped regions."""
    segments: list[str] = []
    current: list[str] = []
    i = 0
    while i < len(command):
        ch = command[i]
        if ch == "\\":
            current.append(command[i : i + 2])
            i += 2
            continue
        if ch in ("'", '"'):
            quote = ch
            j = i + 1
            while j < len(command):
                if command[j] == "\\" and quote == '"':
                    j += 2
                    continue
                if command[j] == quote:
                    j += 1
                    break
                j += 1
            current.append(command[i:j])
            i = j
            continue
        m = pattern.match(command, i)
        if m:
            segments.append("".join(current))
            current = []
            i = m.end()
            continue
        current.append(ch)
        i += 1
    segments.append("".join(current))
    return segments


_CONTROL_RE = re.compile(r"&&|\|\||;|\n|(?<![<>|&])&(?![&>])")
_PIPE_RE = re.compile(r"(?<!>)\|(?!\|)")


def _extract_command_names(command: str) -> list[str]:
    """Extract command names from a shell command string.

    On Windows without bash (PowerShell fallback), uses a simple regex
    since ``shlex.split()`` only handles POSIX quoting.

    Args:
        command: The command string to parse.

    Returns:
        List of command name strings found in the command.
    """
    if sys.platform == "win32" and _WINDOWS_BASH is None:
        return re.findall(r"\b([A-Za-z_][\w.-]*)\b", command)
    names: list[str] = []
    stripped_command = _strip_heredocs(command)
    segments = _split_respecting_quotes(stripped_command, _CONTROL_RE)
    for segment in segments:
        for part in _split_respecting_quotes(segment, _PIPE_RE):
            name = _extract_leading_command_name(part.strip())
            if name:
                names.append(name)
    return names


def _strip_heredocs(command: str) -> str:
    """Strip heredoc content from a bash command.

    Removes everything between << DELIM and DELIM (or <<- DELIM and DELIM),
    so that heredoc body text is not parsed as command arguments.
    """
    return re.sub(
        r"<<-?\s*['\"]?(\w+)['\"]?[^\n]*\n(?:.*?\n)*?[ \t]*\1[ \t]*(?=\n|$)",
        "",
        command,
        flags=re.DOTALL,
    )


def _clean_env() -> dict[str, str]:
    """Return a fresh copy of ``os.environ`` without ``VIRTUAL_ENV``.

    When the agent process runs inside a virtual-env (e.g. the VS Code
    extension's own ``.venv``), the ``VIRTUAL_ENV`` variable leaks into
    child processes and causes ``uv run`` to emit a spurious warning about
    a mismatched environment. Stripping it lets ``uv`` (and other tools)
    discover the correct project ``.venv`` on their own.
    """
    env = os.environ.copy()
    env.pop("VIRTUAL_ENV", None)
    return env


def _format_bash_result(returncode: int, output: str, max_output_chars: int) -> str:
    if returncode != 0:
        msg = f"Error (exit code {returncode}):"
        if output:
            msg += f"\n{output}"
        return _truncate_output(msg, max_output_chars)
    return _truncate_output(output, max_output_chars)


def _kill_process_group(process: subprocess.Popen) -> None:
    """Kill a subprocess and all its children.

    On Windows, uses ``taskkill /T /F`` to kill the entire process tree.
    On Unix, sends ``SIGKILL`` to the process group created by
    ``start_new_session=True``, falling back to ``process.kill()``.

    Args:
        process: The subprocess to terminate.
    """
    if sys.platform == "win32":
        subprocess.run(
            ["taskkill", "/T", "/F", "/PID", str(process.pid)],
            capture_output=True,
        )
    else:
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except OSError:
            try:
                process.kill()
            except OSError:  # pragma: no cover — Popen.send_signal polls first in Python 3.13+
                pass
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:  # pragma: no cover — unreachable after SIGKILL
        pass


def _stop_monitor(
    stop_event: threading.Event,
    process: subprocess.Popen,
    done: threading.Event,
) -> None:
    """Wait for *stop_event* to fire, then kill *process* group.

    Exits when *done* is set (process finished normally) or *stop_event*
    fires (agent was stopped).
    """
    while not done.wait(timeout=0.2):
        if stop_event.is_set():
            _kill_process_group(process)
            return


class UsefulTools:
    """A hardened collection of useful tools with improved security."""

    def __init__(
        self,
        stream_callback: Callable[[str], None] | None = None,
        stop_event: threading.Event | None = None,
    ) -> None:
        self.stream_callback = stream_callback
        self.stop_event = stop_event

    def Read(  # noqa: N802
        self,
        file_path: str,
        max_lines: int = 2000,
    ) -> str:
        """Read file contents.

        Args:
            file_path: Absolute path to file.
            max_lines: Maximum number of lines to return.
        """
        try:
            resolved = Path(file_path).resolve()
            text = resolved.read_text()
            lines = text.splitlines(keepends=True)
            if len(lines) > max_lines:
                return (
                    "".join(lines[:max_lines])
                    + f"\n[truncated: {len(lines) - max_lines} more lines]"
                )
            return text
        except Exception as e:
            _log_exc()
            return f"Error: {e}"

    def Write(  # noqa: N802
        self,
        file_path: str,
        content: str,
    ) -> str:
        """Write content to a file, creating it if it doesn't exist or overwriting if it does.

        Args:
            file_path: Path to the file to write.
            content: The full content to write to the file.
        """
        try:
            resolved = Path(file_path).resolve()
            resolved.parent.mkdir(parents=True, exist_ok=True)
            resolved.write_text(content)
            return f"Successfully wrote {len(content)} characters to {file_path}"
        except Exception as e:
            _log_exc()
            return f"Error: {e}"

    def Edit(  # noqa: N802
        self,
        file_path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
    ) -> str:
        """Performs precise string replacements in files with exact matching.

        Args:
            file_path: Absolute path to the file to modify.
            old_string: Exact text to find and replace.
            new_string: Replacement text, must differ from old_string.
            replace_all: If True, replace all occurrences.

        Returns:
            The output of the edit operation.
        """
        try:
            resolved = Path(file_path).resolve()
            if not resolved.is_file():
                return f"Error: File not found: {file_path}"
            if old_string == new_string:
                return "Error: new_string must be different from old_string"
            content = resolved.read_text()
            count = content.count(old_string)
            if count == 0:
                return "Error: String not found in file"
            if not replace_all and count > 1:
                return (
                    f"Error: String appears {count} times (not unique). "
                    f"Use replace_all=True to replace all occurrences."
                )
            if replace_all:
                new_content = content.replace(old_string, new_string)
            else:
                new_content = content.replace(old_string, new_string, 1)
            resolved.write_text(new_content)
            replaced = count if replace_all else 1
            return f"Successfully replaced {replaced} occurrence(s) in {file_path}"
        except Exception as e:
            _log_exc()
            return f"Error: {e}"

    def Bash(  # noqa: N802
        self,
        command: str,
        description: str,
        timeout_seconds: float = 30,
        max_output_chars: int = 50000,
    ) -> str:
        """Runs a bash command and returns its output.

        Args:
            command: The bash command to run.
            description: A brief description of the command.
            timeout_seconds: Timeout in seconds for the command.
            max_output_chars: Maximum characters in output before truncation.

        Returns:
            The output of the command.
        """
        del description

        disallowed = DISALLOWED_BASH_COMMANDS
        if sys.platform == "win32" and _WINDOWS_BASH is None:
            disallowed = disallowed | DISALLOWED_PS_COMMANDS
        for command_name in _extract_command_names(command):
            if command_name in disallowed:
                return f"Error: Command '{command_name}' is not allowed in Bash tool"

        if self.stream_callback:
            return self._bash_streaming(command, timeout_seconds, max_output_chars)

        env = _clean_env()

        try:
            process = subprocess.Popen(
                **_popen_kwargs(command),
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                env=env,
            )
            done = threading.Event()
            monitor = None
            if self.stop_event:
                monitor = threading.Thread(
                    target=_stop_monitor,
                    args=(self.stop_event, process, done),
                    daemon=True,
                )
                monitor.start()
            try:
                stdout, _ = process.communicate(timeout=timeout_seconds)
            except subprocess.TimeoutExpired:
                _kill_process_group(process)
                try:
                    process.communicate(timeout=5)
                except Exception:  # pragma: no cover — unreachable after SIGKILL
                    pass
                return "Error: Command execution timeout"
            except BaseException:  # pragma: no cover — KeyboardInterrupt timing-dependent
                _kill_process_group(process)
                try:
                    process.communicate(timeout=5)
                except Exception:
                    pass
                raise
            finally:
                done.set()
            return _format_bash_result(process.returncode, stdout, max_output_chars)
        except Exception as e:  # pragma: no cover
            _log_exc()
            return f"Error: {e}"

    def _bash_streaming(self, command: str, timeout_seconds: float, max_output_chars: int) -> str:
        assert self.stream_callback is not None
        process = subprocess.Popen(
            **_popen_kwargs(command),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=_clean_env(),
        )
        timed_out = False
        done = threading.Event()

        def _kill() -> None:
            nonlocal timed_out
            timed_out = True
            _kill_process_group(process)

        timer = threading.Timer(timeout_seconds, _kill)
        timer.start()
        monitor = None
        if self.stop_event:
            monitor = threading.Thread(
                target=_stop_monitor,
                args=(self.stop_event, process, done),
                daemon=True,
            )
            monitor.start()
        try:
            chunks: list[str] = []
            assert process.stdout is not None
            for line in iter(process.stdout.readline, ""):
                chunks.append(line)
                self.stream_callback(line)
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:  # pragma: no cover
                _kill_process_group(process)
        except BaseException:  # pragma: no cover — KeyboardInterrupt timing-dependent
            _kill_process_group(process)
            raise
        finally:
            done.set()
            timer.cancel()
            process.stdout.close()  # type: ignore[union-attr]

        if timed_out:
            return "Error: Command execution timeout"

        output = "".join(chunks)
        return _format_bash_result(process.returncode, output, max_output_chars)
