"""ext.bwrap — local, same-host, bubblewrap namespace isolation (SAND-WP-0013)."""

from __future__ import annotations

import json
import os
import select
import shutil
import signal
import socket
import subprocess
import time
from contextlib import suppress
from pathlib import Path
from typing import Any

from sandboxer.extensions.base import SandboxExtension
from sandboxer.extensions.runtime import RUNTIME_MOUNT, verified_runtime
from sandboxer.models import Profile


class BwrapExtension(SandboxExtension):
    """Provision a local sandbox via bubblewrap (bwrap) kernel namespaces.

    Unlike ext.compose-ssh / ext.vm-packer, this extension never leaves the
    local host: no SSH hop, no container runtime, no remote placement. A new
    user/mount/pid/ipc/uts/net namespace is created per sandbox, kept alive
    by a minimal command broker whose namespace pid is retained for lifecycle
    evidence and teardown. `--unshare-net` with no veth/interface makes
    `network.default: deny` real, rather than declarative-only like the
    other self-hosted extensions.
    """

    def __init__(self, config: dict[str, Any] | None = None) -> None:
        super().__init__(config)
        cfg = self.config
        self.base_dir: str = cfg.get("base_dir", "/tmp/sandboxer-bwrap")
        self.bwrap_bin: str = cfg.get("bwrap_bin", "bwrap")
        self.control_socket_name: str = cfg.get(
            "control_socket_name", ".sandboxer-owner.sock"
        )
        self.ro_binds: list[str] = cfg.get(
            "ro_binds", ["/usr", "/bin", "/lib", "/lib64", "/etc/resolv.conf"]
        )

    def _bwrap_bin(self) -> str:
        return os.environ.get("SANDBOXER_BWRAP_BIN", self.bwrap_bin)

    def _existing_ro_binds(self) -> list[str]:
        return [path for path in self.ro_binds if Path(path).exists()]

    def _bwrap_argv(self, workspace_dir: str, *, info_fd: int | None = None) -> list[str]:
        argv = [
            self._bwrap_bin(),
            "--unshare-user",
            "--unshare-pid",
            "--unshare-ipc",
            "--unshare-uts",
            "--unshare-cgroup",
            "--unshare-net",
            "--tmpfs",
            "/",
            "--proc",
            "/proc",
            "--dev",
            "/dev",
            "--clearenv",
        ]
        if info_fd is not None:
            argv += ["--info-fd", str(info_fd)]
        for path in self._existing_ro_binds():
            argv += ["--ro-bind", path, path]
        runner = Path(__file__).with_name("bwrap_runner.py")
        argv += ["--dir", "/run", "--dir", "/run/sandboxer"]
        argv += ["--ro-bind", str(runner), "/run/sandboxer/bwrap_runner.py"]
        runtime = verified_runtime(self.config)
        if runtime is not None:
            workspace = Path(workspace_dir).resolve()
            if runtime.is_relative_to(workspace) or workspace.is_relative_to(runtime):
                raise ValueError("runtime and workspace must not overlap")
            argv += ["--dir", "/opt", "--dir", "/opt/sandboxer"]
            argv += ["--ro-bind", str(runtime), RUNTIME_MOUNT]
        argv += ["--bind", workspace_dir, workspace_dir]
        argv += ["--chdir", workspace_dir]
        argv += [
            "/usr/bin/python3",
            "/run/sandboxer/bwrap_runner.py",
            workspace_dir,
            f"{workspace_dir}/{self.control_socket_name}",
        ]
        if runtime is not None:
            argv.append("--runtime")
        return argv

    @staticmethod
    def _read_child_pid(proc: subprocess.Popen, info_fd: int) -> int:
        ready, _, _ = select.select([info_fd], [], [], 10)
        if not ready:
            proc.kill()
            raise RuntimeError("timed out waiting for bwrap namespace child pid")
        raw = os.read(info_fd, 16_384)
        try:
            child_pid = int(json.loads(raw)["child-pid"])
        except (json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
            proc.kill()
            raise RuntimeError("bwrap did not report a valid namespace child pid") from exc
        if child_pid <= 0:
            proc.kill()
            raise RuntimeError("bwrap reported an invalid namespace child pid")
        return child_pid

    def provision(
        self, profile: Profile, inputs: dict[str, str], host: str
    ) -> dict[str, str]:
        if profile.network.default != "deny" or profile.network.egress:
            raise ValueError("bwrap currently supports only default-deny with empty egress")
        if profile.setup.secret_refs:
            raise ValueError("bwrap has no setup credential delivery contract")
        runtime = verified_runtime(self.config)
        if runtime is not None and inputs.get("repo"):
            source = Path(inputs["repo"]).resolve()
            if source.is_relative_to(runtime) or runtime.is_relative_to(source):
                raise ValueError("runtime and source checkout must not overlap")
        sandbox_id = self.new_sandbox_id(inputs)
        workspace_dir = f"{self.base_dir}/{sandbox_id}"
        Path(workspace_dir).mkdir(parents=True, exist_ok=True)

        repo = inputs.get("repo")
        if repo:
            repo_path = Path(repo).expanduser().resolve()
            if not repo_path.exists():
                raise FileNotFoundError(f"Repo path does not exist: {repo_path}")
            shutil.copytree(repo_path, workspace_dir, dirs_exist_ok=True)
        Path(workspace_dir).chmod(0o700)

        info_read_fd, info_write_fd = os.pipe()
        try:
            argv = self._bwrap_argv(workspace_dir, info_fd=info_write_fd)
            proc = subprocess.Popen(
                argv,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                start_new_session=True,
                pass_fds=(info_write_fd,),
            )
        finally:
            os.close(info_write_fd)
        try:
            child_pid = self._read_child_pid(proc, info_read_fd)
        finally:
            os.close(info_read_fd)

        return {
            "sandbox_id": sandbox_id,
            "host": host,
            "pid": str(child_pid),
            "supervisor_pid": str(proc.pid),
            "workspace_dir": workspace_dir,
        }

    def wait_ready(self, handle: dict[str, str]) -> dict[str, str]:
        pid = int(handle["pid"])
        if not self._pid_alive(pid):
            raise RuntimeError(f"bwrap process {pid} is not running")
        workspace_dir = handle["workspace_dir"]
        if not Path(workspace_dir).is_dir():
            raise RuntimeError(f"workspace missing: {workspace_dir}")
        control_socket = Path(workspace_dir) / self.control_socket_name
        deadline = time.monotonic() + 5
        while not control_socket.is_socket():
            if not self._pid_alive(pid):
                raise RuntimeError(f"bwrap process {pid} exited before control became ready")
            if time.monotonic() >= deadline:
                raise RuntimeError("bwrap owner control socket did not become ready")
            time.sleep(0.05)
        return {
            "host": handle.get("host", "localhost"),
            "endpoint": f"pid:{pid}",
        }

    def supports_execution(self) -> bool:
        return True

    def _validated_workspace(self, handle: dict[str, str]) -> Path:
        sandbox_id = handle.get("sandbox_id", "")
        if not sandbox_id or "/" in sandbox_id or sandbox_id in {".", ".."}:
            raise RuntimeError("invalid sandbox id in execution handle")
        workspace_value = handle.get("workspace_dir", "")
        if not workspace_value:
            raise RuntimeError("sandbox execution handle has no workspace")
        workspace = Path(workspace_value).resolve(strict=True)
        expected = (Path(self.base_dir).resolve() / sandbox_id).resolve()
        if workspace != expected or not workspace.is_dir():
            raise RuntimeError("refusing execution outside owner-managed sandbox workspace")
        return workspace

    def execute(
        self,
        handle: dict[str, str],
        command: list[str],
        *,
        credential_route_refs: list[str],
        execution_context: dict[str, str],
        timeout_seconds: int,
        max_output_bytes: int,
        stdin_text: str | None = None,
    ) -> dict[str, object]:
        """Ask the broker already inside bwrap to run an argument-vector command."""
        pid = int(handle.get("pid", "0"))
        if pid <= 0 or not self._pid_alive(pid):
            raise RuntimeError(f"bwrap process {pid} is not running")
        workspace = self._validated_workspace(handle)

        request = {
            "command": command,
            "credential_route_refs": credential_route_refs,
            "execution_context": execution_context,
            "timeout_seconds": timeout_seconds,
            "max_output_bytes": max_output_bytes,
            "stdin_text": stdin_text,
        }
        response_limit = max_output_bytes * 2 + 65_536
        chunks: list[bytes] = []
        size = 0
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
            client.settimeout(timeout_seconds + 5)
            client.connect(str(workspace / self.control_socket_name))
            client.sendall(json.dumps(request).encode("utf-8"))
            client.shutdown(socket.SHUT_WR)
            while True:
                chunk = client.recv(65_536)
                if not chunk:
                    break
                size += len(chunk)
                if size > response_limit:
                    raise RuntimeError("bwrap owner response exceeded its declared bound")
                chunks.append(chunk)
        response = json.loads(b"".join(chunks))
        if "boundary_error" in response:
            raise RuntimeError(f"bwrap owner command boundary failed: {response['boundary_error']}")
        return response

    def teardown(self, handle: dict[str, str]) -> dict[str, str]:
        pid_str = handle.get("pid", "")
        killed = False
        if pid_str and self._pid_alive(int(pid_str)):
            pid = int(pid_str)
            try:
                os.killpg(os.getpgid(pid), signal.SIGKILL)
            except (ProcessLookupError, PermissionError):
                with suppress(ProcessLookupError, PermissionError):
                    os.kill(pid, signal.SIGKILL)
            killed = True

        workspace_dir = handle.get("workspace_dir", "")
        removed = False
        if workspace_dir and Path(workspace_dir).exists():
            shutil.rmtree(workspace_dir, ignore_errors=True)
            removed = not Path(workspace_dir).exists()

        return {
            "process_killed": str(killed),
            "workspace_removed": str(removed),
            "workspace_dir": workspace_dir,
        }

    @staticmethod
    def _pid_alive(pid: int) -> bool:
        try:
            os.kill(pid, 0)
        except (ProcessLookupError, PermissionError):
            return False
        return True