from __future__ import annotations

import os
import platform
import shutil
import shlex
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional
from modules.host_sandbox_policy import get_macos_writable_paths


@dataclass
class SandboxPlan:
    command: List[str]
    env: Dict[str, str]
    cwd: Optional[str] = None
    seccomp_bpf_path: Optional[str] = None


class HostSandboxError(RuntimeError):
    pass


def _truthy(name: str, default: str = "1") -> bool:
    return os.environ.get(name, default).strip().lower() not in {"0", "false", "no", "off"}


def host_sandbox_enabled() -> bool:
    return _truthy("HOST_SANDBOX_ENABLED", "1")


def build_host_sandbox_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    system = platform.system()
    if system == "Darwin":
        return _build_macos_plan(command, work_path, env)
    if system == "Linux":
        return _build_linux_plan(command, work_path, env)
    if system == "Windows":
        return _build_windows_plan(command, work_path, env)
    raise HostSandboxError(f"不支持的宿主机系统: {system}")


def build_host_sandbox_readonly_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    system = platform.system()
    if system == "Darwin":
        return _build_macos_readonly_plan(command, work_path, env)
    if system == "Linux":
        return _build_linux_readonly_plan(command, work_path, env)
    if system == "Windows":
        return _build_windows_plan(command, work_path, env)
    raise HostSandboxError(f"不支持的宿主机系统: {system}")

def build_host_sandbox_shell_plan(work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    system = platform.system()
    if system == "Darwin":
        return _build_macos_shell_plan(work_path, env)
    if system == "Linux":
        return _build_linux_shell_plan(work_path, env)
    if system == "Windows":
        return _build_windows_shell_plan(work_path, env)
    raise HostSandboxError(f"不支持的宿主机系统: {system}")


def _build_macos_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    sandbox_exec = shutil.which("sandbox-exec")
    if not sandbox_exec:
        raise HostSandboxError("macOS 未找到 sandbox-exec，拒绝执行宿主机命令。")
    profile = _macos_profile_for_workspace(work_path)
    cmd = [sandbox_exec, "-p", profile, "/bin/bash", "-lc", command]
    return SandboxPlan(command=cmd, env=env, cwd=str(work_path))


def _build_macos_readonly_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    sandbox_exec = shutil.which("sandbox-exec")
    if not sandbox_exec:
        raise HostSandboxError("macOS 未找到 sandbox-exec，拒绝执行宿主机命令。")
    profile = (
        '(version 1) '
        '(deny default) '
        '(allow sysctl-read) '
        '(allow process*) '
        '(allow network*) '
        '(allow file-read*) '
        '(allow file-write* (literal "/dev/null"))'
    )
    cmd = [sandbox_exec, "-p", profile, "/bin/bash", "-lc", command]
    return SandboxPlan(command=cmd, env=env, cwd=str(work_path))

def _build_macos_shell_plan(work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    sandbox_exec = shutil.which("sandbox-exec")
    if not sandbox_exec:
        raise HostSandboxError("macOS 未找到 sandbox-exec，拒绝启动宿主机沙箱终端。")
    profile = _macos_profile_for_workspace(work_path)
    cmd = [sandbox_exec, "-p", profile, "/bin/bash", "-i"]
    return SandboxPlan(command=cmd, env=env, cwd=str(work_path))

def _macos_profile_for_workspace(work_path: Path) -> str:
    workspace = str(work_path.resolve())
    writable_paths = [workspace, "/tmp", "/private/tmp", "/dev/null"]
    for raw in get_macos_writable_paths():
        try:
            expanded = str(Path(raw).expanduser().resolve())
        except Exception:
            continue
        if expanded not in writable_paths:
            writable_paths.append(expanded)
    write_rules: list[str] = []
    for entry in writable_paths:
        if entry == "/dev/null":
            write_rules.append('(literal "/dev/null")')
        else:
            write_rules.append(f'(subpath "{entry}")')
    write_expr = " ".join(write_rules)
    return (
        '(version 1) '
        '(deny default) '
        '(allow sysctl-read) '
        '(allow process*) '
        '(allow network*) '
        '(allow file-read*) '
        f'(allow file-write* {write_expr})'
    )


def _build_linux_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    bwrap = shutil.which("bwrap")
    if not bwrap:
        raise HostSandboxError("Linux 未找到 bubblewrap(bwrap)，拒绝执行宿主机命令。")

    seccomp_bpf = os.environ.get("HOST_SANDBOX_LINUX_SECCOMP_BPF", "").strip()
    if not seccomp_bpf:
        raise HostSandboxError("Linux 缺少 HOST_SANDBOX_LINUX_SECCOMP_BPF，拒绝执行宿主机命令。")
    seccomp_path = Path(seccomp_bpf).expanduser().resolve()
    if not seccomp_path.exists():
        raise HostSandboxError(f"seccomp BPF 文件不存在: {seccomp_path}")
    shell_cmd = ["/bin/bash", "-lc", command]
    return _build_linux_common_plan(work_path, env, shell_cmd, seccomp_path)


def _build_linux_readonly_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    bwrap = shutil.which("bwrap")
    if not bwrap:
        raise HostSandboxError("Linux 未找到 bubblewrap(bwrap)，拒绝执行宿主机命令。")
    seccomp_bpf = os.environ.get("HOST_SANDBOX_LINUX_SECCOMP_BPF", "").strip()
    if not seccomp_bpf:
        raise HostSandboxError("Linux 缺少 HOST_SANDBOX_LINUX_SECCOMP_BPF，拒绝执行宿主机命令。")
    seccomp_path = Path(seccomp_bpf).expanduser().resolve()
    if not seccomp_path.exists():
        raise HostSandboxError(f"seccomp BPF 文件不存在: {seccomp_path}")
    shell_cmd = ["/bin/bash", "-lc", command]
    return _build_linux_common_plan(work_path, env, shell_cmd, seccomp_path, readonly=True)

def _build_linux_shell_plan(work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    bwrap = shutil.which("bwrap")
    if not bwrap:
        raise HostSandboxError("Linux 未找到 bubblewrap(bwrap)，拒绝启动宿主机沙箱终端。")
    seccomp_bpf = os.environ.get("HOST_SANDBOX_LINUX_SECCOMP_BPF", "").strip()
    if not seccomp_bpf:
        raise HostSandboxError("Linux 缺少 HOST_SANDBOX_LINUX_SECCOMP_BPF，拒绝启动宿主机沙箱终端。")
    seccomp_path = Path(seccomp_bpf).expanduser().resolve()
    if not seccomp_path.exists():
        raise HostSandboxError(f"seccomp BPF 文件不存在: {seccomp_path}")
    shell_cmd = ["/bin/bash", "-i"]
    return _build_linux_common_plan(work_path, env, shell_cmd, seccomp_path)

def _build_linux_common_plan(
    work_path: Path,
    env: Dict[str, str],
    shell_cmd: List[str],
    seccomp_path: Path,
    readonly: bool = False,
) -> SandboxPlan:
    bwrap = shutil.which("bwrap")
    if not bwrap:
        raise HostSandboxError("Linux 未找到 bubblewrap(bwrap)。")
    sandbox_root = str(work_path.resolve())
    cmd: List[str] = [
        bwrap,
        "--die-with-parent",
        "--new-session",
        "--unshare-all",
        "--share-net",
        "--ro-bind",
        "/",
        "/",
    ]
    if readonly:
        cmd.extend(["--ro-bind", sandbox_root, sandbox_root])
    else:
        cmd.extend(["--bind", sandbox_root, sandbox_root])
    cmd.extend([
        "--chdir",
        sandbox_root,
        "--proc",
        "/proc",
        "--dev",
        "/dev",
        "--tmpfs",
        "/tmp",
        "--seccomp",
        "__SECCOMP_FD__",
        *shell_cmd,
    ])
    return SandboxPlan(command=cmd, env=env, cwd=sandbox_root, seccomp_bpf_path=str(seccomp_path))


def _build_windows_plan(command: str, work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    wsl = shutil.which("wsl.exe")
    if not wsl:
        raise HostSandboxError("Windows 未找到 wsl.exe（WSL2），拒绝执行宿主机命令。")
    work = shlex.quote(str(work_path))
    shell = f"cd {work} && {command}"
    cmd = [wsl, "bash", "-lc", shell]
    return SandboxPlan(command=cmd, env=env, cwd=str(work_path))

def _build_windows_shell_plan(work_path: Path, env: Dict[str, str]) -> SandboxPlan:
    wsl = shutil.which("wsl.exe")
    if not wsl:
        raise HostSandboxError("Windows 未找到 wsl.exe（WSL2），拒绝启动宿主机沙箱终端。")
    work = shlex.quote(str(work_path))
    shell = f"cd {work} && exec bash -i"
    cmd = [wsl, "bash", "-lc", shell]
    return SandboxPlan(command=cmd, env=env, cwd=str(work_path))