package engine

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
)

// Bwrap runs a RunSpec through the bwrap CLI — filesystem isolation
// only; network, uid mapping target, /proc, /dev, and /tmp come from
// the host by contract.
//
// The invocation is a port of Werkator's BwrapBuildRunner, including
// the parts hardened on a real Hostsharing webspace: bind mountpoints
// are pre-created inside the rootfs (a plain host directory), because
// bwrap cannot mkdir them against the read-only root bind.
type Bwrap struct {
	// Path of the bwrap binary; empty means "bwrap" via PATH.
	Path string
	// Stdio of the sandboxed command; nil fields default to the
	// werkdock process's own.
	Stdout io.Writer
	Stderr io.Writer
	Stdin  io.Reader
}

// DefaultPATH is the PATH inside the sandbox; the environment is
// cleared (docker semantics), so a sane default must be set explicitly.
const DefaultPATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

// Argv assembles the full bwrap command line for spec.
//
// The command is pid 1 of the sandbox (--as-pid-1), as in docker: a
// signal sent to it from outside reaches it if it handles the signal,
// which is what `stop` relies on — bwrap's own init would swallow
// SIGTERM. Zombies of orphaned children are the command's business,
// as they are in a container without --init. bwrap reports the
// command's host pid on fd 3 (--info-fd), which Start reads.
//
// Mount order: the rootfs first; then /proc, /dev, and the tmpfs
// mounts for /tmp and /root, BEFORE the user binds, so a bind whose
// destination lies below them lands inside instead of being shadowed;
// then the user binds in the given order.
func (b *Bwrap) Argv(spec RunSpec) ([]string, error) {
	if spec.RootFS == "" {
		return nil, errors.New("rootfs must be set")
	}
	if !filepath.IsAbs(spec.RootFS) {
		return nil, fmt.Errorf("rootfs must be an absolute path: %s", spec.RootFS)
	}
	if len(spec.Command) == 0 {
		return nil, errors.New("no command specified")
	}
	bin := b.Path
	if bin == "" {
		bin = "bwrap"
	}
	args := []string{
		bin,
		"--unshare-user",
		"--unshare-pid",
		"--as-pid-1",
		"--die-with-parent",
		"--info-fd", "3",
		"--uid", strconv.Itoa(spec.UID),
		"--gid", strconv.Itoa(spec.GID),
		"--ro-bind", spec.RootFS, "/",
		"--proc", "/proc",
		"--dev", "/dev",
		"--tmpfs", "/tmp",
		"--tmpfs", "/root",
		"--tmpfs", "/run",
	}
	for _, m := range spec.Mounts {
		if !filepath.IsAbs(m.Dest) {
			return nil, fmt.Errorf("mount destination must be an absolute path: %s", m.Dest)
		}
		switch m.Mode {
		case MountBind:
			args = append(args, "--bind", m.Source, m.Dest)
		case MountRoBind:
			args = append(args, "--ro-bind", m.Source, m.Dest)
		case MountTmpfs:
			args = append(args, "--tmpfs", m.Dest)
		default:
			return nil, fmt.Errorf("unknown mount mode %d for %s", m.Mode, m.Dest)
		}
	}
	home := spec.Home
	if home == "" {
		home = "/root"
	}
	args = append(args, "--clearenv")
	for _, e := range environment(home, spec.Env) {
		args = append(args, "--setenv", e.Key, e.Value)
	}
	workdir := spec.Workdir
	if workdir == "" {
		workdir = "/"
	}
	if WorkdirWritable(spec) {
		// docker creates a missing workdir; bwrap can only do so where
		// the sandbox is writable (--dir is mkdir -p, a no-op if present).
		args = append(args, "--dir", workdir)
	}
	args = append(args, "--chdir", workdir, "--")
	args = append(args, spec.Command...)
	return args, nil
}

// environment is HOME and PATH with werkdock's defaults, then the
// spec's variables in order; a later value for the same key replaces
// the earlier one in place, so each key is set exactly once.
func environment(home string, env []EnvVar) []EnvVar {
	all := []EnvVar{{Key: "HOME", Value: home}, {Key: "PATH", Value: DefaultPATH}}
	for _, e := range env {
		replaced := false
		for i := range all {
			if all[i].Key == e.Key {
				all[i].Value, replaced = e.Value, true
				break
			}
		}
		if !replaced {
			all = append(all, e)
		}
	}
	return all
}

// implicitMounts are the tmpfs mounts every sandbox carries, in front of
// the user mounts.
// /run holds runtime state by FHS definition (pid files, sockets, locks)
// and is /var/run's target in practically every image; services write
// there without declaring a VOLUME (RFC 0005).
var implicitMounts = []Mount{{Mode: MountTmpfs, Dest: "/tmp"}, {Mode: MountTmpfs, Dest: "/root"}, {Mode: MountTmpfs, Dest: "/run"}}

// coveringMount returns the mount whose destination is the longest
// prefix of path — the one the path is actually served from, since
// later and deeper mounts shadow earlier ones. ok is false when the
// path lies in the read-only rootfs itself.
func coveringMount(spec RunSpec, path string) (m Mount, ok bool) {
	best := -1
	for _, cand := range append(append([]Mount{}, implicitMounts...), spec.Mounts...) {
		dest := filepath.Clean(cand.Dest)
		if path == dest || strings.HasPrefix(path, strings.TrimSuffix(dest, "/")+"/") {
			if len(dest) >= best {
				best, m, ok = len(dest), cand, true
			}
		}
	}
	return m, ok
}

// WorkdirWritable reports whether spec.Workdir lies in a writable area
// of the sandbox: a tmpfs or a read-write bind.
func WorkdirWritable(spec RunSpec) bool {
	if spec.Workdir == "" {
		return false
	}
	m, ok := coveringMount(spec, filepath.Clean(spec.Workdir))
	return ok && m.Mode != MountRoBind
}

// CheckWorkdir verifies that a workdir outside the writable areas exists
// where the sandbox will see it — in the rootfs or the read-only bind
// covering it — so a typo fails with a clear message instead of bwrap's
// "Can't chdir".
func CheckWorkdir(spec RunSpec) error {
	if spec.Workdir == "" || WorkdirWritable(spec) {
		return nil
	}
	workdir := filepath.Clean(spec.Workdir)
	hostPath := ""
	if m, ok := coveringMount(spec, workdir); ok {
		hostPath = filepath.Join(m.Source, strings.TrimPrefix(workdir, filepath.Clean(m.Dest)))
	} else {
		var err error
		if hostPath, err = rootfsPath(spec.RootFS, workdir); err != nil {
			return err
		}
	}
	if fi, err := os.Stat(hostPath); err != nil || !fi.IsDir() {
		return fmt.Errorf("working directory %s does not exist in the image and lies outside the sandbox's writable areas (tmpfs and read-write binds)", spec.Workdir)
	}
	return nil
}

// EnsureMountpoints pre-creates the mountpoints of spec inside the
// rootfs directory. bwrap creates mountpoints against the sandbox view,
// which is the read-only rootfs bind — every destination missing from
// the rootfs fails with "Read-only file system". The rootfs directory
// itself is a plain host directory, so the mountpoints are created
// there; bwrap then finds them and has nothing left to mkdir.
//
// Anything that already exists in the rootfs is left alone (e.g.
// /etc/resolv.conf is a file many rootfs archives ship). A bind whose
// source is a regular file gets a file mountpoint, not a directory.
func EnsureMountpoints(spec RunSpec) error {
	for _, dest := range []string{"/proc", "/dev", "/tmp", "/root"} {
		if err := ensureDir(spec.RootFS, dest); err != nil {
			return err
		}
	}
	for _, m := range spec.Mounts {
		target, err := rootfsPath(spec.RootFS, m.Dest)
		if err != nil {
			return err
		}
		if _, err := os.Lstat(target); err == nil {
			continue
		}
		if m.Mode == MountTmpfs {
			if err := os.MkdirAll(target, 0o755); err != nil {
				return err
			}
			continue
		}
		src, err := os.Stat(m.Source)
		if err != nil {
			return fmt.Errorf("bind source %s: %w", m.Source, err)
		}
		if src.Mode().IsRegular() {
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				return err
			}
			f, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o644)
			if err != nil {
				return err
			}
			if err := f.Close(); err != nil {
				return err
			}
			continue
		}
		if err := os.MkdirAll(target, 0o755); err != nil {
			return err
		}
	}
	return nil
}

func ensureDir(rootfs, dest string) error {
	target, err := rootfsPath(rootfs, dest)
	if err != nil {
		return err
	}
	if _, statErr := os.Lstat(target); statErr == nil {
		return nil
	}
	return os.MkdirAll(target, 0o755)
}

// rootfsPath resolves dest inside rootfs and refuses destinations that
// escape it — werkdock assembles mounts from user input, so this must
// hold even for hostile paths.
func rootfsPath(rootfs, dest string) (string, error) {
	root := filepath.Clean(rootfs)
	target := filepath.Join(root, dest)
	prefix := root
	if !strings.HasSuffix(prefix, string(filepath.Separator)) {
		prefix += string(filepath.Separator)
	}
	if target != root && !strings.HasPrefix(target, prefix) {
		return "", fmt.Errorf("bind destination escapes the rootfs: %s", dest)
	}
	return target, nil
}

// process is a running bwrap invocation.
type process struct {
	cmd      *exec.Cmd
	childPID int
}

func (p *process) ChildPID() int { return p.childPID }

func (p *process) Signal(sig os.Signal) error {
	s, ok := sig.(syscall.Signal)
	if !ok {
		return fmt.Errorf("unsupported signal %v", sig)
	}
	return syscall.Kill(p.childPID, s)
}

// Wait returns the command's exit code; bwrap propagates the child's
// code (128+signal for a signalled child), so it passes through.
func (p *process) Wait() (int, error) {
	err := p.cmd.Wait()
	if err == nil {
		return 0, nil
	}
	var exit *exec.ExitError
	if errors.As(err, &exit) {
		if exit.ExitCode() >= 0 {
			return exit.ExitCode(), nil
		}
		if ws, ok := exit.Sys().(syscall.WaitStatus); ok && ws.Signaled() {
			return 128 + int(ws.Signal()), nil
		}
	}
	return 0, err
}

// Start launches spec and returns once bwrap has reported the
// sandboxed command's host pid. bwrap runs in a process group of its
// own, so a terminal's Ctrl-C reaches werkdock, which forwards it as
// the supervisor sees fit, rather than killing bwrap outright.
func (b *Bwrap) Start(spec RunSpec) (Process, error) {
	argv, err := b.Argv(spec)
	if err != nil {
		return nil, err
	}
	if err := EnsureMountpoints(spec); err != nil {
		return nil, err
	}
	if err := CheckWorkdir(spec); err != nil {
		return nil, err
	}
	infoR, infoW, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	defer infoR.Close()
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Stdout = b.Stdout
	if cmd.Stdout == nil {
		cmd.Stdout = os.Stdout
	}
	cmd.Stderr = b.Stderr
	if cmd.Stderr == nil {
		cmd.Stderr = os.Stderr
	}
	cmd.Stdin = b.Stdin
	cmd.ExtraFiles = []*os.File{infoW} // fd 3, see Argv
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		infoW.Close()
		return nil, err
	}
	infoW.Close()
	var info struct {
		ChildPID int `json:"child-pid"`
	}
	if err := json.NewDecoder(infoR).Decode(&info); err != nil || info.ChildPID <= 0 {
		// bwrap failed before the sandbox existed; its message went
		// to stderr, its exit code is the only thing left to report.
		code, werr := (&process{cmd: cmd}).Wait()
		if werr != nil {
			return nil, werr
		}
		return nil, fmt.Errorf("bwrap exited with code %d before starting the command", code)
	}
	return &process{cmd: cmd, childPID: info.ChildPID}, nil
}

// Run executes spec and returns the command's exit code; bwrap
// propagates the child's code, so the caller can pass it through.
func (b *Bwrap) Run(spec RunSpec) (int, error) {
	p, err := b.Start(spec)
	if err != nil {
		return 0, err
	}
	return p.Wait()
}