// Copyright 2023 The Shac Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package engine

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"slices"
	"strings"

	"go.fuchsia.dev/shac-project/shac/internal/execsupport"
	"go.fuchsia.dev/shac-project/shac/internal/sandbox"
	"go.starlark.net/starlark"
)

// subprocess represents an in-progress subprocess as returned by ctx.os.exec().
type subprocess struct {
	cmd            *exec.Cmd
	args           []string
	stdout         *bytes.Buffer
	stderr         *bytes.Buffer
	raiseOnFailure bool
	okRetcodes     []int
	tempDir        string
	errs           <-chan error

	waitCalled bool
}

var _ starlark.HasAttrs = (*subprocess)(nil)

func (s *subprocess) String() string {
	return fmt.Sprintf("<subprocess %q>", strings.Join(s.args, " "))
}

func (s *subprocess) Type() string {
	return "subprocess"
}

func (s *subprocess) Truth() starlark.Bool {
	return true
}

func (s *subprocess) Freeze() {
}

func (s *subprocess) Hash() (uint32, error) {
	return 0, errors.New("unhashable type: subprocess")
}

func (s *subprocess) Attr(name string) (starlark.Value, error) {
	switch name {
	case "wait":
		return subprocessWaitBuiltin.BindReceiver(s), nil
	default:
		return nil, nil
	}
}

func (s *subprocess) AttrNames() []string {
	return []string{"wait"}
}

func (s *subprocess) wait() (starlark.Value, error) {
	if s.waitCalled {
		return nil, fmt.Errorf("wait was already called")
	}
	s.waitCalled = true
	val, err := s.waitInner()
	if err2 := s.cleanup(); err == nil {
		err = err2
	}
	return val, err
}

func (s *subprocess) waitInner() (starlark.Value, error) {
	retcode := 0
	if err := <-s.errs; err != nil {
		var errExit *exec.ExitError
		if errors.As(err, &errExit) {
			retcode = errExit.ExitCode()
		} else {
			// Something other than a normal non-zero exit.
			return nil, err
		}
	}

	// Limits output to 10Mib. If it needs more, a file should probably be used.
	// If there is a use case, it's fine to increase.
	const limit = 10 * 1024 * 1024
	if s.stdout.Len() > limit {
		return nil, errors.New("process returned too much stdout")
	}
	if s.stderr.Len() > limit {
		return nil, errors.New("process returned too much stderr")
	}

	if !slices.Contains(s.okRetcodes, retcode) && s.raiseOnFailure {
		var msgBuilder strings.Builder
		msgBuilder.WriteString(fmt.Sprintf("command failed with exit code %d: %s", retcode, s.args))
		if s.stderr.Len() > 0 {
			msgBuilder.WriteString("\n")
			msgBuilder.WriteString(s.stderr.String())
		}
		return nil, errors.New(msgBuilder.String())
	}
	return toValue("completed_subprocess", starlark.StringDict{
		"retcode": starlark.MakeInt(retcode),
		"stdout":  starlark.String(s.stdout.String()),
		"stderr":  starlark.String(s.stderr.String()),
	}), nil
}

func (s *subprocess) cleanup() error {
	// Wait for the subprocess to launch before trying to kill it. s.startErrs
	// gets closed after the subprocess starts, so even if the error has already
	// been received by `wait()`, this receive will return due to the channel
	// being closed.
	<-s.errs
	// Kill the process before doing any other cleanup steps to ensure resources
	// are no longer in use before cleaning them up.
	var err error
	// If Process is not nil then the command successfully started. If
	// ProcessState is nil then the command hasn't yet completed.
	if s.cmd.Process != nil && s.cmd.ProcessState == nil {
		err = s.cmd.Process.Kill()
		// Kill() is non-blocking, so it's necessary to wait for the process to
		// exit before cleaning up resources.
		_ = s.cmd.Wait()
	}

	if err2 := os.RemoveAll(s.tempDir); err == nil {
		err = err2
	}

	buffers.push(s.stdout)
	buffers.push(s.stderr)
	s.stdout, s.stderr = nil, nil

	return err
}

var subprocessWaitBuiltin = newBoundBuiltin("wait", func(ctx context.Context, s *shacState, name string, self starlark.Value, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	if err := starlark.UnpackArgs(name, args, kwargs); err != nil {
		return nil, err
	}
	return self.(*subprocess).wait()
})

// ctxOsExec implements the native function ctx.os.exec().
//
// Make sure to update //doc/stdlib.star whenever this function is modified.
func ctxOsExec(ctx context.Context, s *shacState, name string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var argcmd starlark.Sequence
	var argcwd starlark.String
	var argenv = starlark.NewDict(0)
	var argstdin starlark.Value = starlark.None
	var argraiseOnFailure starlark.Bool = true
	var argallowNetwork starlark.Bool
	var argokRetcodes starlark.Value = starlark.None
	if err := starlark.UnpackArgs(name, args, kwargs,
		"cmd", &argcmd,
		"cwd?", &argcwd,
		"env?", &argenv,
		"stdin?", &argstdin,
		"allow_network?", &argallowNetwork,
		"ok_retcodes?", &argokRetcodes,
		"raise_on_failure?", &argraiseOnFailure,
	); err != nil {
		return nil, err
	}
	if argcmd.Len() == 0 {
		return nil, errors.New("cmdline must not be an empty list")
	}

	var okRetcodes []int
	if argokRetcodes == starlark.None {
		okRetcodes = append(okRetcodes, 0)
	} else {
		if !argraiseOnFailure {
			return nil, fmt.Errorf("cannot combine \"ok_retcodes\" and \"raise_on_failure=False\"")
		}
		seqOkRetcodes, ok := argokRetcodes.(starlark.Sequence)
		if ok {
			okRetcodes = sequenceToInts(seqOkRetcodes)
		}
		if !ok || okRetcodes == nil {
			return nil, fmt.Errorf("for parameter \"ok_retcodes\": got %s, wanted sequence of ints", argokRetcodes)
		}
	}

	var cleanupFuncs []func() error
	defer func() {
		for _, f := range cleanupFuncs {
			// Ignore errors during cleanup because cleanupFuncs will only be
			// populated if another error occurred prior to starting the
			// subprocess.
			_ = f()
		}
	}()

	tempDir, err := s.newTempDir()
	if err != nil {
		return nil, err
	}

	cleanupFuncs = append(cleanupFuncs, func() error {
		return os.RemoveAll(tempDir)
	})

	env := map[string]string{
		"PATH":    os.Getenv("PATH"),
		"TEMP":    tempDir,
		"TMPDIR":  tempDir,
		"TEMPDIR": tempDir,
	}
	if runtime.GOROOT() != "" {
		// TODO(olivernewman): This is necessary because checks for shac itself
		// assume Go is pre-installed. Switch to a hermetic Go installation that
		// installs Go in the checkout directory, and stop explicitly mounting
		// $GOROOT and adding it to $PATH.
		env["PATH"] = strings.Join([]string{
			filepath.Join(runtime.GOROOT(), "bin"),
			env["PATH"],
		}, string(os.PathListSeparator))
	}

	var passthroughMounts []sandbox.Mount
	for _, pte := range s.passthroughEnv {
		val, ok := os.LookupEnv(pte.Name)
		if !ok {
			continue
		}
		env[pte.Name] = val
		if pte.IsPath {
			passthroughMounts = append(passthroughMounts, sandbox.Mount{
				Path:     val,
				Writable: pte.Writeable,
			})
		}
	}

	for _, item := range argenv.Items() {
		k, ok := item[0].(starlark.String)
		if !ok {
			return nil, fmt.Errorf("\"env\" key is not a string: %s", item[0])
		}
		// TODO(olivernewman): This is unnecessarily strict - commands should
		// not set $PATH in `env`, but we should allow prepending to $PATH with
		// `env_prefixes`, and add an option to not inherit the value of $PATH
		// if better hermeticity is desired.
		if k == "PATH" {
			return nil, fmt.Errorf("$PATH cannot be overridden")
		}
		v, ok := item[1].(starlark.String)
		if !ok {
			return nil, fmt.Errorf("\"env\" value is not a string: %s", item[1])
		}
		env[string(k)] = string(v)
	}

	var stdin io.Reader
	switch s := argstdin.(type) {
	case starlark.String:
		stdin = strings.NewReader(string(s))
	case starlark.Bytes:
		stdin = bytes.NewReader([]byte(s))
	case starlark.NoneType:
	default:
		return nil, fmt.Errorf("for parameter \"stdin\": got %s, want str or bytes", argstdin.Type())
	}

	cwd := filepath.Join(s.root, s.subdir)
	if s := string(argcwd); s != "" {
		cwd, err = absPath(s, cwd)
		if err != nil {
			return nil, err
		}
	}

	fullCmd := sequenceToStrings(argcmd)
	if fullCmd == nil {
		return nil, fmt.Errorf("for parameter \"cmd\": got %s, want sequence of str", argcmd.Type())
	}

	if filepath.IsAbs(fullCmd[0]) {
		// Stat to make sure the entrypoint executable exists rather than
		// letting nsjail fail, for consistency with the non-absolute path case.
		if _, err = os.Stat(fullCmd[0]); err != nil {
			return nil, err
		}
	} else {
		// nsjail doesn't do $PATH-based resolution of the command it's given.
		// Do this resolution unconditionally for consistency across platforms
		// even though it's not necessary when not using nsjail. This also
		// ensures that relative file paths are interpreted relative to the root
		// directory, rather than the directory from which shac is run.
		absPath := filepath.Join(s.root, s.subdir, fullCmd[0])
		if _, err = os.Stat(absPath); err != nil {
			// exec.LookPath() doesn't do $PATH-based lookup for paths
			// containing slashes, so no point in trying it if the file doesn't
			// exist.
			if !errors.Is(err, os.ErrNotExist) || strings.Contains(fullCmd[0], "/") {
				return nil, err
			}
			// If the path doesn't exist in the root, fall back to a $PATH
			// lookup.
			absPath, err = exec.LookPath(fullCmd[0])
			if err != nil {
				return nil, err
			}
		}
		fullCmd[0] = absPath
	}

	config := &sandbox.Config{
		Cmd:          fullCmd,
		Cwd:          cwd,
		AllowNetwork: bool(argallowNetwork),
		Env:          env,
	}
	// config.Mounts is ignored for the moment on Windows.
	if runtime.GOOS != "windows" {
		config.Mounts = []sandbox.Mount{
			// TODO(olivernewman): Mount the checkout read-only unconditionally.
			{Path: s.root, Writable: s.writableRoot},
			// OS-provided utilities.
			{Path: "/dev/null", Writable: true},
			{Path: "/dev/urandom"},
			{Path: "/dev/zero"},
			// DNS configs.
			{Path: "/etc/nsswitch.conf"},
			{Path: "/etc/resolv.conf"},
			// Required for https.
			{Path: "/etc/ssl/certs"},
			// These are required for bash to work.
			{Path: "/lib"},
			{Path: "/lib64"},
			// OS header files.
			{Path: "/usr/include"},
			// System compilers.
			{Path: "/usr/lib"},
			{Path: "/usr/share"},
			// Make the parent directory of tempDir available, since it is the root
			// of all ctx.os.tempdir() calls, which can be used as scratch pads for
			// this executable.
			{Path: filepath.Dir(tempDir), Writable: true},
		}
		config.Mounts = append(config.Mounts, passthroughMounts...)

		// Explicitly mount standard binary directories. On systems with a
		// merged /usr layout (like Debian 12+), /bin is a symlink to /usr/bin.
		// A restricted $PATH might omit /bin, so we must ensure it's mounted
		// for scripts that rely on standard shebangs like #!/bin/sh to function
		// correctly.
		for _, p := range []string{"/bin", "/usr/bin", "/sbin", "/usr/sbin"} {
			if fi, errStat := os.Stat(p); errStat == nil && fi.IsDir() {
				config.Mounts = append(config.Mounts, sandbox.Mount{Path: p})
			}
		}

		// TODO(olivernewman): This is necessary because checks for shac itself
		// assume Go is pre-installed. Switch to a hermetic Go installation that
		// installs Go in the checkout directory, and stop explicitly mounting
		// $GOROOT and adding it to $PATH.
		if runtime.GOROOT() != "" {
			config.Mounts = append(config.Mounts, sandbox.Mount{Path: runtime.GOROOT()})
		}

		// Mount all directories listed in $PATH.
		for p := range strings.SplitSeq(env["PATH"], string(os.PathListSeparator)) {
			// $PATH may contain invalid elements. Filter them out.
			if !filepath.IsAbs(p) {
				// Relative paths in $PATH are not allowed.
				continue
			}
			var fi os.FileInfo
			if fi, err = os.Stat(p); err != nil || !fi.IsDir() {
				// Skip $PATH elements that don't exist or point to
				// non-directories.
				continue
			}
			config.Mounts = append(config.Mounts, sandbox.Mount{Path: p})
		}
	}

	cmd := s.sandbox.Command(ctx, config)

	stdout, stderr := buffers.get(), buffers.get()
	// TODO(olivernewman): Also handle commands that may output non-utf-8 bytes.
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	cmd.Stdin = stdin

	errs := make(chan error, 1)
	// Run the command in a non-blocking goroutine so exec() calls don't block
	// if there's already the maximum number of subprocesses running. wait()
	// will block until the subprocess starts *and* finishes.
	go func() {
		errs <- func() error {
			if err := s.subprocessSem.Acquire(ctx, 1); err != nil {
				return err
			}
			defer s.subprocessSem.Release(1)
			log.Printf("Running command: %s", cmd)
			return execsupport.Run(cmd)
		}()
		// Signals to subprocess.wait() that the subprocess is done, whether or
		// not it was successful.
		close(errs)
	}()

	proc := &subprocess{
		cmd:            cmd,
		args:           sequenceToStrings(argcmd),
		stdout:         stdout,
		stderr:         stderr,
		raiseOnFailure: bool(argraiseOnFailure),
		okRetcodes:     okRetcodes,
		tempDir:        tempDir,
		errs:           errs,
	}

	// Only clean up now if starting the subprocess failed; otherwise it will
	// get cleaned up by wait().
	cleanupFuncs = cleanupFuncs[:0]

	chk := ctxCheck(ctx)
	chk.subprocesses = append(chk.subprocesses, proc)
	return proc, nil
}

// sequenceToInts converts a starlark sequence (list, tuple) into a slice of
// ints.
func sequenceToInts(s starlark.Sequence) []int {
	out := make([]int, 0, s.Len())
	iter := s.Iterate()
	var x starlark.Value
	for iter.Next(&x) {
		i, ok := x.(starlark.Int)
		if !ok {
			return nil
		}
		i64, ok := i.Int64()
		if !ok {
			return nil
		}
		if i64 > math.MaxInt || i64 < math.MinInt {
			return nil
		}
		out = append(out, int(i64))
	}
	return out
}