"""Safe shell execution tool with allowlist enforcement."""
from __future__ import annotations

import subprocess

from core.state import PlanStep, StepResult, StepStatus

# Commands explicitly allowed
ALLOWED_COMMANDS: frozenset[str] = frozenset({"pwd", "ls", "echo", "cat", "python", "git"})

# Commands explicitly blocked (belt-and-suspenders)
BLOCKED_COMMANDS: frozenset[str] = frozenset({
    "bash", "sh", "rm", "sudo", "curl", "wget",
    "chmod", "chown", "docker", "systemctl",
})

# For commands with restricted sub-commands
RESTRICTED_SUBCOMMANDS: dict[str, set[str]] = {
    "python": {"--version"},
    "git": {"status"},
}


def run_shell(step: PlanStep) -> StepResult:
    """Execute a whitelisted shell command via subprocess (shell=False)."""
    raw_command: str = step.payload.get("command", "").strip()
    parts = raw_command.split()

    if not parts:
        return StepResult(step_id=step.step_id, status=StepStatus.FAILED, error="Empty command")

    binary = parts[0]

    if binary in BLOCKED_COMMANDS:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error=f"Command '{binary}' is blocked by security policy",
        )

    if binary not in ALLOWED_COMMANDS:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error=f"Command '{binary}' is not in allowlist. Allowed: {sorted(ALLOWED_COMMANDS)}",
        )

    # Validate restricted sub-commands
    if binary in RESTRICTED_SUBCOMMANDS:
        allowed_sub = RESTRICTED_SUBCOMMANDS[binary]
        sub = " ".join(parts[1:])
        if sub not in allowed_sub:
            return StepResult(
                step_id=step.step_id,
                status=StepStatus.FAILED,
                error=f"Only '{binary} {sorted(allowed_sub)}' is permitted, got: '{binary} {sub}'",
            )

    try:
        proc = subprocess.run(
            parts,
            shell=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
        output = proc.stdout + proc.stderr
        if proc.returncode != 0:
            return StepResult(
                step_id=step.step_id,
                status=StepStatus.FAILED,
                output_text=output,
                error=f"Non-zero exit code: {proc.returncode}",
            )
        return StepResult(step_id=step.step_id, status=StepStatus.DONE, output_text=output)

    except subprocess.TimeoutExpired:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error="Command timed out after 10s",
        )
    except FileNotFoundError:
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.FAILED,
            error=f"Binary '{binary}' not found on PATH",
        )