"""Bash command execution tool.

Pure function executor for bash commands that accepts ToolCall and returns ToolResult.
"""

import os
import subprocess
from pathlib import Path

import trio

from wafer_core.rollouts.dtypes import (
    Tool,
    ToolCall,
    ToolFunction,
    ToolFunctionParameter,
    ToolResult,
)

# Constants
MAX_OUTPUT_SIZE = 30_000  # 30KB (matches Claude Code's default)


# ── Tool Definition ──────────────────────────────────────────────────────────

BASH_TOOL = Tool(
    type="function",
    function=ToolFunction(
        name="bash",
        description="Execute a bash command in the current working directory. Returns stdout and stderr.",
        parameters=ToolFunctionParameter(
            type="object",
            properties={
                "command": {"type": "string", "description": "Bash command to execute"},
                "timeout": {"type": "integer", "description": "Timeout in seconds (default: 120)"},
            },
        ),
        required=["command"],
    ),
)


# ── Helper Functions ─────────────────────────────────────────────────────────


def _split_compound_command(command: str) -> list[str]:
    """Split a compound bash command into individual commands.

    Handles &&, ||, ;, and | operators while respecting quotes.
    Returns list of individual command strings.

    Examples:
        "cd foo && ls" -> ["cd foo", "ls"]
        "echo 'a && b'" -> ["echo 'a && b'"]  # quoted && preserved
        "a | b | c" -> ["a", "b", "c"]
    """
    commands = []
    current = []
    in_single_quote = False
    in_double_quote = False
    i = 0

    while i < len(command):
        char = command[i]

        # Handle quotes
        if char == "'" and not in_double_quote:
            in_single_quote = not in_single_quote
            current.append(char)
        elif char == '"' and not in_single_quote:
            in_double_quote = not in_double_quote
            current.append(char)
        # Handle operators (only outside quotes)
        elif not in_single_quote and not in_double_quote:
            # Check for && or ||
            if i + 1 < len(command) and command[i : i + 2] in ("&&", "||"):
                if current:
                    commands.append("".join(current).strip())
                    current = []
                i += 2
                continue
            # Check for ; or |
            if char in (";", "|"):
                if current:
                    commands.append("".join(current).strip())
                    current = []
            else:
                current.append(char)
        else:
            current.append(char)

        i += 1

    # Add final command
    if current:
        cmd = "".join(current).strip()
        if cmd:
            commands.append(cmd)

    return commands


def check_bash_allowlist(command: str, allowlist: list[str] | None) -> str | None:
    """Check if all commands in a (potentially compound) bash command are allowed.

    Splits compound commands (using &&, ||, ;, |) and checks each individual
    command against the allowlist.

    Returns None if all commands allowed, or an error message if any blocked.
    """
    if allowlist is None:
        return None  # No allowlist = all commands allowed

    # Split compound command into individual commands
    individual_commands = _split_compound_command(command)

    if not individual_commands:
        # Empty command, allow it (will fail at execution anyway)
        return None

    # Check each command against the allowlist
    blocked_commands = []
    for cmd in individual_commands:
        allowed = False
        for prefix in allowlist:
            if cmd.startswith(prefix):
                allowed = True
                break
        if not allowed:
            blocked_commands.append(cmd)

    if blocked_commands:
        allowed_str = ", ".join(f"'{p}'" for p in allowlist)
        blocked_str = ", ".join(f"'{c}'" for c in blocked_commands)
        return (
            f"Command(s) not allowed: {blocked_str}\n"
            f"This environment only permits commands starting with: {allowed_str}"
        )

    return None  # All commands allowed


# ── Pure Function Executor ───────────────────────────────────────────────────


async def exec_bash(
    tool_call: ToolCall,
    working_dir: Path,
    cancel_scope: trio.CancelScope | None = None,
    bash_allowlist: list[str] | None = None,
) -> ToolResult:
    """Execute bash command (pure function - takes working_dir as parameter).

    Args:
        tool_call: The tool call containing command and timeout.
        working_dir: Working directory for command execution.
        cancel_scope: Optional trio cancel scope.
        bash_allowlist: Optional list of allowed command prefixes (from template).
    """
    command = tool_call.args["command"]
    timeout = tool_call.args.get("timeout", 120)

    # Check bash allowlist if set
    if error := check_bash_allowlist(command, bash_allowlist):
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=error,
        )

    try:
        # Inherit current process environment (includes venv PATH)
        env = os.environ.copy()
        result = await trio.to_thread.run_sync(
            lambda: subprocess.run(
                ["sh", "-c", command],
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=str(working_dir),
                env=env,
            ),
        )

        output = ""
        if result.stdout:
            output += result.stdout
        if result.stderr:
            if output:
                output += "\n"
            output += result.stderr

        # Truncate if too large
        if len(output) > MAX_OUTPUT_SIZE:
            output = output[:MAX_OUTPUT_SIZE] + "\n\n... (output truncated)"

        if result.returncode != 0:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content=output or "(no output)",
                error=f"Command exited with code {result.returncode}",
            )

        return ToolResult(
            tool_call_id=tool_call.id, is_error=False, content=output or "(no output)"
        )

    except subprocess.TimeoutExpired:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Command timed out after {timeout} seconds",
        )
    except trio.Cancelled:
        return ToolResult(
            tool_call_id=tool_call.id, is_error=True, content="", error="Command aborted"
        )
