// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gosh

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

var (
	errAlreadyCalledStart = errors.New("gosh: already called Cmd.Start")
	errAlreadyCalledWait  = errors.New("gosh: already called Cmd.Wait")
	errCloseStdout        = errors.New("gosh: use NopWriteCloser(os.Stdout) to prevent stdout from being closed")
	errCloseStderr        = errors.New("gosh: use NopWriteCloser(os.Stderr) to prevent stderr from being closed")
	errDidNotCallStart    = errors.New("gosh: did not call Cmd.Start")
	errProcessExited      = errors.New("gosh: process exited")
)

// Cmd represents a command. Not thread-safe.
// Public fields should not be modified after calling Start.
type Cmd struct {
	// Err is the most recent error from this Cmd (may be nil).
	Err error
	// Path is the path of the command to run.
	Path string
	// Vars is the map of env vars for this Cmd.
	Vars map[string]string
	// Args is the list of args for this Cmd, starting with the resolved path.
	// Note, we set Args[0] to the resolved path (rather than the user-specified
	// name) so that a command started by Shell can reliably determine the path to
	// its executable.
	Args []string
	// IgnoreParentExit, if true, makes it so the child process does not exit when
	// its parent exits. Only takes effect if the child process was spawned via
	// Shell.Fn or Shell.Main, or explicitly calls InitChildMain.
	IgnoreParentExit bool
	// ExitAfter, if non-zero, specifies that the child process should exit after
	// the given duration has elapsed. Only takes effect if the child process was
	// spawned via Shell.Fn or Shell.Main, or explicitly calls InitChildMain.
	ExitAfter time.Duration
	// PropagateOutput is inherited from Shell.Opts.PropagateChildOutput.
	PropagateOutput bool
	// OutputDir is inherited from Shell.Opts.ChildOutputDir.
	OutputDir string
	// ExitErrorIsOk specifies whether an *exec.ExitError should be reported via
	// Shell.HandleError.
	ExitErrorIsOk bool
	// Stdin is a string to write to the child's stdin.
	Stdin string
	// Internal state.
	sh               *Shell
	c                *exec.Cmd
	stdinWriteCloser io.WriteCloser // from exec.Cmd.StdinPipe
	calledStart      bool
	calledWait       bool
	cond             *sync.Cond
	waitChan         chan error
	started          bool // protected by sh.cleanupMu
	exited           bool // protected by cond.L
	stdoutWriters    []io.Writer
	stderrWriters    []io.Writer
	closers          []io.Closer
	recvReady        bool              // protected by cond.L
	recvVars         map[string]string // protected by cond.L
}

// Clone returns a new Cmd with a copy of this Cmd's configuration.
func (c *Cmd) Clone() *Cmd {
	c.sh.Ok()
	res, err := c.clone()
	c.handleError(err)
	return res
}

// StdinPipe returns a thread-safe WriteCloser backed by a buffered pipe for the
// command's stdin. The returned pipe will be closed when the process exits, but
// may also be closed earlier by the caller, e.g. if the command does not exit
// until its stdin is closed. Must be called before Start. It is safe to call
// StdinPipe multiple times; calls after the first return the pipe created by
// the first call.
func (c *Cmd) StdinPipe() io.WriteCloser {
	c.sh.Ok()
	res, err := c.stdinPipe()
	c.handleError(err)
	return res
}

// StdoutPipe returns a Reader backed by a buffered pipe for the command's
// stdout. Must be called before Start. May be called more than once; each
// invocation creates a new pipe.
func (c *Cmd) StdoutPipe() io.Reader {
	c.sh.Ok()
	res, err := c.stdoutPipe()
	c.handleError(err)
	return res
}

// StderrPipe returns a Reader backed by a buffered pipe for the command's
// stderr. Must be called before Start. May be called more than once; each
// invocation creates a new pipe.
func (c *Cmd) StderrPipe() io.Reader {
	c.sh.Ok()
	res, err := c.stderrPipe()
	c.handleError(err)
	return res
}

// AddStdoutWriter configures this Cmd to tee the child's stdout to the given
// WriteCloser, which will be closed when the process exits.
//
// If the same WriteCloser is passed to both AddStdoutWriter and
// AddStderrWriter, Cmd will ensure that its methods are never called
// concurrently and that Close is only called once.
//
// Use NopWriteCloser to extend a Writer to a WriteCloser, or to prevent an
// existing WriteCloser from being closed. It is an error to pass in os.Stdout
// or os.Stderr, since they shouldn't be closed.
func (c *Cmd) AddStdoutWriter(wc io.WriteCloser) {
	c.sh.Ok()
	c.handleError(c.addStdoutWriter(wc))
}

// AddStderrWriter configures this Cmd to tee the child's stderr to the given
// WriteCloser, which will be closed when the process exits.
//
// If the same WriteCloser is passed to both AddStdoutWriter and
// AddStderrWriter, Cmd will ensure that its methods are never called
// concurrently and that Close is only called once.
//
// Use NopWriteCloser to extend a Writer to a WriteCloser, or to prevent an
// existing WriteCloser from being closed. It is an error to pass in os.Stdout
// or os.Stderr, since they shouldn't be closed.
func (c *Cmd) AddStderrWriter(wc io.WriteCloser) {
	c.sh.Ok()
	c.handleError(c.addStderrWriter(wc))
}

// Start starts the command.
func (c *Cmd) Start() {
	c.sh.Ok()
	c.handleError(c.start())
}

// AwaitReady waits for the child process to call SendReady. Must not be called
// before Start or after Wait.
func (c *Cmd) AwaitReady() {
	c.sh.Ok()
	c.handleError(c.awaitReady())
}

// AwaitVars waits for the child process to send values for the given vars
// (using SendVars). Must not be called before Start or after Wait.
func (c *Cmd) AwaitVars(keys ...string) map[string]string {
	c.sh.Ok()
	res, err := c.awaitVars(keys...)
	c.handleError(err)
	return res
}

// Wait waits for the command to exit.
func (c *Cmd) Wait() {
	c.sh.Ok()
	c.handleError(c.wait())
}

// Kill causes the process to exit immediately.
func (c *Cmd) Kill() {
	c.sh.Ok()
	c.handleError(c.kill())
}

// Signal sends a signal to the process.
func (c *Cmd) Signal(sig os.Signal) {
	c.sh.Ok()
	c.handleError(c.signal(sig))
}

// Shutdown sends a signal to the process, then waits for it to exit.
func (c *Cmd) Shutdown(sig os.Signal) {
	c.sh.Ok()
	c.handleError(c.shutdown(sig))
}

// Run calls Start followed by Wait.
func (c *Cmd) Run() {
	c.sh.Ok()
	c.handleError(c.run())
}

// Stdout calls Start followed by Wait, then returns the command's stdout.
func (c *Cmd) Stdout() string {
	c.sh.Ok()
	res, err := c.stdout()
	c.handleError(err)
	return res
}

// StdoutStderr calls Start followed by Wait, then returns the command's stdout
// and stderr.
func (c *Cmd) StdoutStderr() (string, string) {
	c.sh.Ok()
	stdout, stderr, err := c.stdoutStderr()
	c.handleError(err)
	return stdout, stderr
}

// CombinedOutput calls Start followed by Wait, then returns the command's
// combined stdout and stderr.
func (c *Cmd) CombinedOutput() string {
	c.sh.Ok()
	res, err := c.combinedOutput()
	c.handleError(err)
	return res
}

// Pid returns the command's PID, or -1 if the command has not been started.
func (c *Cmd) Pid() int {
	if !c.started {
		return -1
	}
	return c.c.Process.Pid
}

////////////////////////////////////////
// Internals

func newCmdInternal(sh *Shell, vars map[string]string, path string, args []string) (*Cmd, error) {
	c := &Cmd{
		Path:     path,
		Vars:     vars,
		Args:     append([]string{path}, args...),
		sh:       sh,
		c:        &exec.Cmd{},
		cond:     sync.NewCond(&sync.Mutex{}),
		waitChan: make(chan error, 1),
		recvVars: map[string]string{},
	}
	// Protect against concurrent signal-triggered Shell.cleanup().
	sh.cleanupMu.Lock()
	defer sh.cleanupMu.Unlock()
	if sh.calledCleanup {
		return nil, errAlreadyCalledCleanup
	}
	sh.cmds = append(sh.cmds, c)
	return c, nil
}

func newCmd(sh *Shell, vars map[string]string, name string, args ...string) (*Cmd, error) {
	// Mimics https://golang.org/src/os/exec/exec.go Command.
	if filepath.Base(name) == name {
		if lp, err := exec.LookPath(name); err != nil {
			return nil, err
		} else {
			name = lp
		}
	}
	return newCmdInternal(sh, vars, name, args)
}

func (c *Cmd) errorIsOk(err error) bool {
	if c.ExitErrorIsOk {
		if _, ok := err.(*exec.ExitError); ok {
			return true
		}
	}
	return err == nil
}

func (c *Cmd) handleError(err error) {
	c.Err = err
	if !c.errorIsOk(err) {
		c.sh.HandleError(err)
	}
}

func (c *Cmd) closeClosers() {
	// If the same WriteCloser was passed to both AddStdoutWriter and
	// AddStderrWriter, we should only close it once.
	cm := map[io.Closer]bool{}
	for _, c := range c.closers {
		if !cm[c] {
			cm[c] = true
			c.Close() // best-effort; ignore returned error
		}
	}
}

func (c *Cmd) isRunning() bool {
	if !c.started {
		return false
	}
	c.cond.L.Lock()
	defer c.cond.L.Unlock()
	return !c.exited
}

// recvWriter listens for gosh messages from a child process.
type recvWriter struct {
	c          *Cmd
	buf        bytes.Buffer
	readPrefix bool // if true, we've read len(msgPrefix) for the current line
	skipLine   bool // if true, ignore bytes until next '\n'
}

func (w *recvWriter) Write(p []byte) (n int, err error) {
	for _, b := range p {
		if b == '\n' {
			if w.readPrefix && !w.skipLine {
				m := msg{}
				if err := json.Unmarshal(w.buf.Bytes(), &m); err != nil {
					return 0, err
				}
				switch m.Type {
				case typeReady:
					w.c.cond.L.Lock()
					w.c.recvReady = true
					w.c.cond.Signal()
					w.c.cond.L.Unlock()
				case typeVars:
					w.c.cond.L.Lock()
					w.c.recvVars = mergeMaps(w.c.recvVars, m.Vars)
					w.c.cond.Signal()
					w.c.cond.L.Unlock()
				default:
					return 0, fmt.Errorf("unknown message type: %q", m.Type)
				}
			}
			// Reset state for next line.
			w.readPrefix, w.skipLine = false, false
			w.buf.Reset()
		} else if !w.skipLine {
			w.buf.WriteByte(b)
			if !w.readPrefix && w.buf.Len() == len(msgPrefix) {
				w.readPrefix = true
				prefix := string(w.buf.Next(len(msgPrefix)))
				if prefix != msgPrefix {
					w.skipLine = true
				}
			}
		}
	}
	return len(p), nil
}

type lockedWriter struct {
	mu *sync.Mutex
	w  io.Writer
}

func (w lockedWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	n, err := w.w.Write(p)
	w.mu.Unlock()
	return n, err
}

func (c *Cmd) makeStdoutStderr() (io.Writer, io.Writer, error) {
	c.stdoutWriters = append(c.stdoutWriters, &recvWriter{c: c})
	if c.PropagateOutput {
		c.stdoutWriters = append(c.stdoutWriters, os.Stdout)
		c.stderrWriters = append(c.stderrWriters, os.Stderr)
	}
	if c.OutputDir != "" {
		t := time.Now().Format("20060102.150405.000000")
		name := filepath.Join(c.OutputDir, filepath.Base(c.Path)+"."+t)
		const flags = os.O_WRONLY | os.O_CREATE | os.O_EXCL
		switch file, err := os.OpenFile(name+".stdout", flags, 0600); {
		case err != nil:
			return nil, nil, err
		default:
			c.stdoutWriters = append(c.stdoutWriters, file)
			c.closers = append(c.closers, file)
		}
		switch file, err := os.OpenFile(name+".stderr", flags, 0600); {
		case err != nil:
			return nil, nil, err
		default:
			c.stderrWriters = append(c.stderrWriters, file)
			c.closers = append(c.closers, file)
		}
	}
	switch hasOut, hasErr := len(c.stdoutWriters) > 0, len(c.stderrWriters) > 0; {
	case hasOut && hasErr:
		// Make writes synchronous between stdout and stderr. This ensures all
		// writers that capture both will see the same ordering, and don't need to
		// worry about concurrent writes.
		sharedMu := &sync.Mutex{}
		stdout := lockedWriter{sharedMu, io.MultiWriter(c.stdoutWriters...)}
		stderr := lockedWriter{sharedMu, io.MultiWriter(c.stderrWriters...)}
		return stdout, stderr, nil
	case hasOut:
		return io.MultiWriter(c.stdoutWriters...), nil, nil
	case hasErr:
		return nil, io.MultiWriter(c.stderrWriters...), nil
	}
	return nil, nil, nil
}

func (c *Cmd) clone() (*Cmd, error) {
	vars := copyMap(c.Vars)
	args := make([]string, len(c.Args))
	copy(args, c.Args)
	res, err := newCmdInternal(c.sh, vars, c.Path, args[1:])
	if err != nil {
		return nil, err
	}
	res.IgnoreParentExit = c.IgnoreParentExit
	res.ExitAfter = c.ExitAfter
	res.PropagateOutput = c.PropagateOutput
	res.OutputDir = c.OutputDir
	res.ExitErrorIsOk = c.ExitErrorIsOk
	res.Stdin = c.Stdin
	return res, nil
}

func (c *Cmd) stdinPipe() (io.WriteCloser, error) {
	if c.calledStart {
		return nil, errAlreadyCalledStart
	}
	if c.stdinWriteCloser != nil {
		return c.stdinWriteCloser, nil
	}
	var err error
	c.stdinWriteCloser, err = c.c.StdinPipe()
	return c.stdinWriteCloser, err
}

func (c *Cmd) stdoutPipe() (io.Reader, error) {
	if c.calledStart {
		return nil, errAlreadyCalledStart
	}
	p := NewBufferedPipe()
	c.stdoutWriters = append(c.stdoutWriters, p)
	c.closers = append(c.closers, p)
	return p, nil
}

func (c *Cmd) stderrPipe() (io.Reader, error) {
	if c.calledStart {
		return nil, errAlreadyCalledStart
	}
	p := NewBufferedPipe()
	c.stderrWriters = append(c.stderrWriters, p)
	c.closers = append(c.closers, p)
	return p, nil
}

func (c *Cmd) addStdoutWriter(wc io.WriteCloser) error {
	switch {
	case c.calledStart:
		return errAlreadyCalledStart
	case wc == os.Stdout:
		return errCloseStdout
	case wc == os.Stderr:
		return errCloseStderr
	}
	c.stdoutWriters = append(c.stdoutWriters, wc)
	c.closers = append(c.closers, wc)
	return nil
}

func (c *Cmd) addStderrWriter(wc io.WriteCloser) error {
	switch {
	case c.calledStart:
		return errAlreadyCalledStart
	case wc == os.Stdout:
		return errCloseStdout
	case wc == os.Stderr:
		return errCloseStderr
	}
	c.stderrWriters = append(c.stderrWriters, wc)
	c.closers = append(c.closers, wc)
	return nil
}

// TODO(sadovsky): Maybe wrap every child process with a "supervisor" process
// that calls InitChildMain.

func (c *Cmd) start() error {
	if c.calledStart {
		return errAlreadyCalledStart
	}
	c.calledStart = true
	// Protect against Cmd.start() writing to c.c.Process concurrently with
	// signal-triggered Shell.cleanup() reading from it.
	c.sh.cleanupMu.Lock()
	defer c.sh.cleanupMu.Unlock()
	if c.sh.calledCleanup {
		return errAlreadyCalledCleanup
	}
	// Configure the command.
	c.c.Path = c.Path
	vars := copyMap(c.Vars)
	if c.IgnoreParentExit {
		delete(vars, envWatchParent)
	} else {
		vars[envWatchParent] = "1"
	}
	if c.ExitAfter == 0 {
		delete(vars, envExitAfter)
	} else {
		vars[envExitAfter] = c.ExitAfter.String()
	}
	c.c.Env = mapToSlice(vars)
	c.c.Args = c.Args
	if c.Stdin != "" {
		if c.stdinWriteCloser != nil {
			return errors.New("gosh: cannot both set Stdin and call StdinPipe")
		}
		c.c.Stdin = strings.NewReader(c.Stdin)
	}
	var err error
	if c.c.Stdout, c.c.Stderr, err = c.makeStdoutStderr(); err != nil {
		return err
	}
	// Start the command.
	onExit := func(err error) {
		c.cond.L.Lock()
		c.exited = true
		c.cond.Signal()
		c.cond.L.Unlock()
		c.closeClosers()
		c.waitChan <- err
	}
	if err = c.c.Start(); err != nil {
		onExit(errors.New("gosh: start failed"))
		return err
	}
	c.started = true
	// Spawn a "waiter" goroutine that calls exec.Cmd.Wait and thus waits for the
	// process to exit. Calling exec.Cmd.Wait here rather than in gosh.Cmd.Wait
	// ensures that the child process is reaped once it exits. Note, gosh.Cmd.wait
	// blocks on waitChan.
	go func() {
		onExit(c.c.Wait())
	}()
	return nil
}

// TODO(sadovsky): Maybe add optional timeouts for
// Cmd.{awaitReady,awaitVars,wait}.

func (c *Cmd) awaitReady() error {
	if !c.started {
		return errDidNotCallStart
	} else if c.calledWait {
		return errAlreadyCalledWait
	}
	c.cond.L.Lock()
	defer c.cond.L.Unlock()
	for !c.exited && !c.recvReady {
		c.cond.Wait()
	}
	// Return nil error if both conditions triggered simultaneously.
	if !c.recvReady {
		return errProcessExited
	}
	return nil
}

func (c *Cmd) awaitVars(keys ...string) (map[string]string, error) {
	if !c.started {
		return nil, errDidNotCallStart
	} else if c.calledWait {
		return nil, errAlreadyCalledWait
	}
	wantKeys := map[string]bool{}
	for _, key := range keys {
		wantKeys[key] = true
	}
	res := map[string]string{}
	updateRes := func() {
		for k, v := range c.recvVars {
			if _, ok := wantKeys[k]; ok {
				res[k] = v
			}
		}
	}
	c.cond.L.Lock()
	defer c.cond.L.Unlock()
	updateRes()
	for !c.exited && len(res) < len(wantKeys) {
		c.cond.Wait()
		updateRes()
	}
	// Return nil error if both conditions triggered simultaneously.
	if len(res) < len(wantKeys) {
		return nil, errProcessExited
	}
	return res, nil
}

func (c *Cmd) wait() error {
	if !c.started {
		return errDidNotCallStart
	} else if c.calledWait {
		return errAlreadyCalledWait
	}
	c.calledWait = true
	return <-c.waitChan
}

// Note: We check for this particular error message to handle the unavoidable
// race between sending a signal to a process and the process exiting.
// https://golang.org/src/os/exec_unix.go
// https://golang.org/src/os/exec_windows.go
const errFinished = "os: process already finished"

func (c *Cmd) kill() error {
	if !c.started {
		return errDidNotCallStart
	}
	if !c.isRunning() {
		return nil
	}
	if err := c.c.Process.Kill(); err != nil && err.Error() != errFinished {
		return err
	}
	return nil
}

func (c *Cmd) signal(sig os.Signal) error {
	if !c.started {
		return errDidNotCallStart
	}
	if !c.isRunning() {
		return nil
	}
	if err := c.c.Process.Signal(sig); err != nil && err.Error() != errFinished {
		return err
	}
	return nil
}

func (c *Cmd) shutdown(sig os.Signal) error {
	if err := c.signal(sig); err != nil {
		return err
	}
	if err := c.wait(); err != nil {
		if _, ok := err.(*exec.ExitError); !ok {
			return err
		}
	}
	return nil
}

func (c *Cmd) run() error {
	if err := c.start(); err != nil {
		return err
	}
	return c.wait()
}

func (c *Cmd) stdout() (string, error) {
	if c.calledStart {
		return "", errAlreadyCalledStart
	}
	var stdout bytes.Buffer
	c.stdoutWriters = append(c.stdoutWriters, &stdout)
	err := c.run()
	return stdout.String(), err
}

func (c *Cmd) stdoutStderr() (string, string, error) {
	if c.calledStart {
		return "", "", errAlreadyCalledStart
	}
	var stdout, stderr bytes.Buffer
	c.stdoutWriters = append(c.stdoutWriters, &stdout)
	c.stderrWriters = append(c.stderrWriters, &stderr)
	err := c.run()
	return stdout.String(), stderr.String(), err
}

func (c *Cmd) combinedOutput() (string, error) {
	if c.calledStart {
		return "", errAlreadyCalledStart
	}
	var output bytes.Buffer
	c.stdoutWriters = append(c.stdoutWriters, &output)
	c.stderrWriters = append(c.stderrWriters, &output)
	err := c.run()
	return output.String(), err
}