from __future__ import annotations

import asyncio
import os
import shlex
import time
from dataclasses import dataclass
from shutil import which
from pathlib import Path
from typing import Any

from ..schemas import SecurityMode
from ..util import truncate_text

BLOCKED_SAFE_TOKENS = {
    "rm",
    "mv",
    "chmod",
    "chown",
    "dd",
    "mkfs",
    "mount",
    "umount",
    "iptables",
    "kill",
    "killall",
    "pkill",
    "shutdown",
    "reboot",
    "poweroff",
    "curl",
    "wget",
    "ssh",
    "scp",
    "nc",
    "ncat",
    "su",
    "sudo",
}
BLOCKED_SAFE_PATTERNS = ["&&", "||", ";", ">", "<", "$(", "`", "|", "\n"]
SAFE_SUBCOMMANDS = {
    "pip": {"list", "show", "freeze"},
    "pip3": {"list", "show", "freeze"},
    "python": {"--version", "-V"},
    "python3": {"--version", "-V"},
}


class ShellSafetyError(RuntimeError):
    pass


@dataclass(slots=True)
class ShellRunResult:
    command: str
    cwd: str
    mode: SecurityMode
    exit_code: int
    stdout: str
    stderr: str
    duration_seconds: float
    timed_out: bool = False

    def to_dict(self) -> dict[str, Any]:
        return {
            "command": self.command,
            "cwd": self.cwd,
            "mode": self.mode,
            "exit_code": self.exit_code,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "duration_seconds": round(self.duration_seconds, 3),
            "timed_out": self.timed_out,
        }


class ShellExecutor:
    def __init__(self, config: dict[str, Any]):
        self.config = config
        shell = config.get("shell", {}) or {}
        timeout_value = shell.get("timeout_seconds", 60)
        self.timeout_seconds = None if timeout_value in (None, 0, "0", False) else int(timeout_value)
        output_limit = shell.get("max_output_chars", 12000)
        self.max_output_chars = None if output_limit in (None, 0, "0", False) else int(output_limit)
        self.safe_allowed_commands = set(shell.get("safe_allowed_commands", []))
        self.project_root = Path(config.get("project_root", ".")).resolve()
        self.workspace_dir = Path(config.get("workspace_dir", self.project_root / "workspace")).resolve()
        self.allowed_roots = [
            self.project_root,
            self.workspace_dir,
            Path(config.get("data_dir", self.project_root / "data")).resolve(),
            Path(config.get("logs_dir", self.project_root / "logs")).resolve(),
            Path(config.get("backups_dir", self.project_root / "backups")).resolve(),
        ]

    def resolve_cwd(self, requested: str | None, mode: SecurityMode) -> Path:
        if not requested:
            return Path.cwd() if mode == "god" else self.workspace_dir
        candidate = Path(requested).expanduser()
        if not candidate.is_absolute():
            candidate = (self.workspace_dir / candidate).resolve()
        else:
            candidate = candidate.resolve()
        if mode == "god":
            return candidate
        if any(candidate == root or candidate.is_relative_to(root) for root in self.allowed_roots):
            return candidate
        raise ShellSafetyError(f"cwd outside allowed roots for mode={mode}: {candidate}")

    def validate(self, command: str, mode: SecurityMode) -> None:
        if mode != "safe":
            return
        if not command.strip():
            raise ShellSafetyError("empty command")
        if any(pattern in command for pattern in BLOCKED_SAFE_PATTERNS):
            raise ShellSafetyError("unsafe shell metacharacters are blocked in safe mode")
        argv = shlex.split(command, posix=True)
        if not argv:
            raise ShellSafetyError("empty command")
        executable = Path(argv[0]).name
        if executable not in self.safe_allowed_commands:
            raise ShellSafetyError(f"command not allowed in safe mode: {executable}")
        if executable in SAFE_SUBCOMMANDS:
            if len(argv) < 2 or argv[1] not in SAFE_SUBCOMMANDS[executable]:
                raise ShellSafetyError(f"subcommand not allowed in safe mode for {executable}")
        blocked_tokens = {Path(token).name for token in argv[1:]}
        if BLOCKED_SAFE_TOKENS.intersection(blocked_tokens):
            raise ShellSafetyError("dangerous token detected in safe mode")

    async def run(
        self,
        command: str,
        *,
        mode: SecurityMode,
        cwd: str | None = None,
        timeout_seconds: int | None = None,
        extra_env: dict[str, str] | None = None,
    ) -> ShellRunResult:
        self.validate(command, mode)
        resolved_cwd = self.resolve_cwd(cwd, mode)
        timeout = timeout_seconds if timeout_seconds not in (None, 0, "0", False) else self.timeout_seconds
        if mode == "god":
            timeout = None
        env = os.environ.copy()
        env.update(extra_env or {})
        env["OPENCLAW_LITE_MODE"] = mode

        started = time.perf_counter()
        process = await asyncio.create_subprocess_shell(
            command,
            cwd=str(resolved_cwd),
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        timed_out = False
        try:
            if timeout is None:
                stdout_bytes, stderr_bytes = await process.communicate()
            else:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(process.communicate(), timeout=timeout)
        except TimeoutError:
            timed_out = True
            process.kill()
            stdout_bytes, stderr_bytes = await process.communicate()
        duration = time.perf_counter() - started
        stdout_raw = stdout_bytes.decode("utf-8", errors="replace")
        stderr_raw = stderr_bytes.decode("utf-8", errors="replace")
        max_chars = None if mode == "god" else self.max_output_chars
        stdout = truncate_text(stdout_raw, max_chars) if max_chars else stdout_raw
        stderr = truncate_text(stderr_raw, max_chars) if max_chars else stderr_raw
        return ShellRunResult(
            command=command,
            cwd=str(resolved_cwd),
            mode=mode,
            exit_code=process.returncode if process.returncode is not None else 124,
            stdout=stdout,
            stderr=stderr,
            duration_seconds=duration,
            timed_out=timed_out,
        )
