package main

// Job control: in an interactive shell every pipeline runs in its own process
// group, which gets the terminal while it runs in the foreground. Ctrl+Z stops
// the foreground job, "&" starts a job in the background, and the builtins
// jobs, fg, bg and wait manage the jobs. Without a terminal (scripts, gosh -c)
// there is no job control, but "&" and wait still work.

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/chzyer/readline"
	"golang.org/x/sys/unix"
)

// proc is a stage of a job: a process, or a builtin or a stage that failed to
// start (both without process).
type proc struct {
	cmd      *exec.Cmd  // nil if the stage is no process
	done     chan error // the result of a stage without process
	finished bool
	stopped  bool
	err      error
}

// job is a pipeline that was started by the shell.
type job struct {
	id    int    // number in the job table, 0 if not in the table
	pgid  int    // process group, 0 without job control or processes
	text  string // the command line
	procs []*proc
	tty   *readline.State // terminal modes of the job when it was stopped
}

var (
	jobControl       bool            // job control is enabled
	interactiveShell bool            // the shell reads commands from a user
	ttyFd            int             // the terminal of the shell
	shellPgid        int             // the process group of the shell
	origPgid         int             // the foreground process group at start
	shellTTY         *readline.State // terminal modes of the shell
	jobTable         []*job          // background and stopped jobs
	exitWarned       bool            // "there are stopped jobs" was shown
)

// initJobControl enables job control if stdin is a terminal: the shell moves
// into its own process group and takes the terminal.
func initJobControl() {
	fd := int(os.Stdin.Fd())
	if !readline.IsTerminal(fd) {
		return
	}
	pgrp, err := unix.IoctlGetInt(fd, unix.TIOCGPGRP)
	if err != nil {
		return
	}
	pid := os.Getpid()
	if syscall.Getpgrp() != pid {
		if err := syscall.Setpgid(0, 0); err != nil {
			return
		}
	}
	origPgid, ttyFd, shellPgid = pgrp, fd, pid
	if err := setForeground(pid); err != nil {
		return
	}
	// Ctrl+Z must not stop the shell itself. (signal.Ignore would be
	// inherited by the child processes.)
	signal.Notify(make(chan os.Signal, 1), syscall.SIGTSTP)
	jobControl = true
}

// releaseTerminal gives the terminal back to the process group that had it
// when the shell started.
func releaseTerminal() {
	if jobControl && origPgid > 0 && origPgid != shellPgid {
		setForeground(origPgid)
	}
}

// setForeground makes pgid the foreground process group of the terminal.
func setForeground(pgid int) error {
	// The shell gets SIGTTOU if it does this while it is in the background
	// (when it takes the terminal back from a job), unless the signal is
	// ignored. It is only ignored for this call, since ignored signals are
	// inherited by child processes.
	signal.Ignore(syscall.SIGTTOU)
	defer signal.Reset(syscall.SIGTTOU)
	return unix.IoctlSetPointerInt(ttyFd, unix.TIOCSPGRP, pgid)
}

// ttyIndex returns the index of the file in fds that is the terminal of the
// shell, or -1.
func ttyIndex(fds [3]*os.File) int {
	tty, err := os.Stdin.Stat()
	if err != nil {
		return -1
	}
	for i, f := range fds {
		if info, err := f.Stat(); err == nil && os.SameFile(info, tty) {
			return i
		}
	}
	return -1
}

// finishedProc returns a proc without process that already has its result.
func finishedProc(err error) *proc {
	p := &proc{done: make(chan error, 1)}
	p.done <- err
	return p
}

// running reports whether a process of the job is running.
func (j *job) running() bool {
	for _, p := range j.procs {
		if p.cmd != nil && !p.finished && !p.stopped {
			return true
		}
	}
	return false
}

// stopped reports whether the job is stopped: no process is running, but
// one is stopped.
func (j *job) stopped() bool {
	if j.running() {
		return false
	}
	for _, p := range j.procs {
		if p.stopped {
			return true
		}
	}
	return false
}

// finished reports whether all stages of the job have finished.
func (j *job) finished() bool {
	for _, p := range j.procs {
		if !p.finished {
			return false
		}
	}
	return true
}

// update records a status change of process pid.
func (j *job) update(pid int, ws syscall.WaitStatus) {
	for _, p := range j.procs {
		if p.cmd == nil || p.cmd.Process.Pid != pid {
			continue
		}
		switch {
		case ws.Stopped():
			p.stopped = true
		case ws.Continued():
			p.stopped = false
		default:
			p.finished, p.stopped = true, false
			p.err = waitStatusError(ws)
			p.cmd.Process.Release()
		}
	}
}

// waitStatusError converts the status of a finished process into an error:
// nil for exit code 0, otherwise a statusError with the exit code, or 128+n
// if the process was killed by signal n.
func waitStatusError(ws syscall.WaitStatus) error {
	switch {
	case ws.Signaled():
		return statusError(128 + int(ws.Signal()))
	case ws.ExitStatus() != 0:
		return statusError(ws.ExitStatus())
	}
	return nil
}

// wait waits until no process of the job is running any more (all finished
// or, with job control, stopped). With nohang it only collects the status
// changes that are already available and does not block.
func (j *job) wait(nohang bool) {
	flags := 0
	if jobControl && j.pgid != 0 {
		flags |= syscall.WUNTRACED | syscall.WCONTINUED
	}
	if nohang {
		flags |= syscall.WNOHANG
	}

	for {
		// With job control all processes of the job are in its process
		// group, otherwise they are waited for one by one.
		target := 0
		if jobControl && j.pgid != 0 {
			if !j.running() && !(nohang && j.stopped()) {
				break
			}
			target = -j.pgid
		} else {
			if nohang {
				// Check every process once
				for _, p := range j.procs {
					if p.cmd == nil || p.finished {
						continue
					}
					var ws syscall.WaitStatus
					if pid, err := syscall.Wait4(p.cmd.Process.Pid, &ws, flags, nil); err == nil && pid > 0 {
						j.update(pid, ws)
					}
				}
				break
			}
			for _, p := range j.procs {
				if p.cmd != nil && !p.finished {
					target = p.cmd.Process.Pid
					break
				}
			}
			if target == 0 {
				break
			}
		}

		var ws syscall.WaitStatus
		pid, err := syscall.Wait4(target, &ws, flags, nil)
		if errors.Is(err, syscall.EINTR) {
			continue
		}
		if err != nil {
			// Should not happen; do not wait forever
			for _, p := range j.procs {
				if p.cmd != nil && !p.finished {
					p.finished, p.stopped = true, false
				}
			}
			break
		}
		if pid == 0 {
			break // nohang and nothing changed
		}
		j.update(pid, ws)
	}

	// Stages without process (builtins in pipelines)
	for _, p := range j.procs {
		if p.cmd != nil || p.finished {
			continue
		}
		if nohang || j.stopped() {
			select {
			case p.err = <-p.done:
				p.finished = true
			default:
			}
		} else {
			p.err = <-p.done
			p.finished = true
		}
	}
}

// result returns the error of the last stage of a finished job (so $? is its
// exit status, like in sh) and prints the errors of the other stages.
func (j *job) result() error {
	for _, p := range j.procs[:len(j.procs)-1] {
		printError(p.err)
	}
	return j.procs[len(j.procs)-1].err
}

// waitForeground waits for a job in the foreground. If it is stopped, it is
// added to the job table.
func waitForeground(j *job) error {
	j.wait(false)
	if jobControl && j.pgid != 0 {
		setForeground(shellPgid)
		if j.stopped() {
			j.tty, _ = readline.GetState(ttyFd)
		}
		if shellTTY != nil {
			readline.Restore(ttyFd, shellTTY)
		}
	}

	if j.stopped() {
		addJob(j)
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, j.describe("Stopped"))
		return statusError(128 + int(syscall.SIGTSTP))
	}
	removeJob(j)
	j.wait(false) // wait for the builtins, if any
	return j.result()
}

// addJob adds j to the job table (if it is not there yet) and makes it the
// current job.
func addJob(j *job) {
	removeJob(j)
	if j.id == 0 {
		for _, other := range jobTable {
			j.id = max(j.id, other.id)
		}
		j.id++
	}
	jobTable = append(jobTable, j)
}

// removeJob removes j from the job table.
func removeJob(j *job) {
	for i, other := range jobTable {
		if other == j {
			jobTable = append(jobTable[:i], jobTable[i+1:]...)
			return
		}
	}
}

// describe formats a job for jobs and the notifications, like bash:
// "[1]+  Stopped                 sleep 30". The current job (the last one in
// the table, i.e. the one stopped or put into the background most recently) is
// marked with +, the previous one with -.
func (j *job) describe(state string) string {
	mark := ' '
	if n := len(jobTable); n > 0 && jobTable[n-1] == j {
		mark = '+'
	} else if n > 1 && jobTable[n-2] == j {
		mark = '-'
	}
	text := j.text
	if state == "Running" {
		text += " &"
	}
	return fmt.Sprintf("[%d]%c  %-24s%s", j.id, mark, state, sanitize(text))
}

// state returns the state of a job for jobs and the notifications.
func (j *job) state() string {
	switch {
	case j.stopped():
		return "Stopped"
	case !j.finished():
		return "Running"
	}
	if status := exitStatus(j.procs[len(j.procs)-1].err); status != 0 {
		return "Exit " + strconv.Itoa(status)
	}
	return "Done"
}

// hasStoppedJobs reports whether a job in the table is stopped.
func hasStoppedJobs() bool {
	for _, j := range jobTable {
		if j.stopped() {
			return true
		}
	}
	return false
}

// reapJobs collects the status changes of the background jobs, reports
// finished and newly stopped jobs (if notify is set) and removes finished
// jobs from the table.
func reapJobs(notify bool) {
	for _, j := range append([]*job(nil), jobTable...) {
		wasStopped := j.stopped()
		j.wait(true)
		switch {
		case j.finished():
			if notify {
				for _, p := range j.procs {
					printError(p.err)
				}
				fmt.Fprintln(os.Stderr, j.describe(j.state()))
			}
			removeJob(j)
		case j.stopped() && !wasStopped:
			addJob(j) // the current job now
			if notify {
				fmt.Fprintln(os.Stderr, j.describe("Stopped"))
			}
		}
	}
}

// findJob returns the job for a job spec: %N or N for job N, %+, %% or nothing
// for the current job, %- for the previous one.
func findJob(cmd string, args []string) (*job, error) {
	spec := "%+"
	if len(args) > 1 {
		spec = args[1]
	}
	n := len(jobTable)
	switch spec {
	case "%+", "%%", "%":
		if n > 0 {
			return jobTable[n-1], nil
		}
		return nil, fmt.Errorf("%s: current: no such job", cmd)
	case "%-":
		if n > 1 {
			return jobTable[n-2], nil
		}
		return nil, fmt.Errorf("%s: previous: no such job", cmd)
	}
	id, err := strconv.Atoi(strings.TrimPrefix(spec, "%"))
	if err == nil {
		for _, j := range jobTable {
			if j.id == id {
				return j, nil
			}
		}
	}
	return nil, fmt.Errorf("%s: %s: no such job", cmd, sanitize(spec))
}

// continueJob sends SIGCONT to a stopped job.
func continueJob(j *job) {
	for _, p := range j.procs {
		p.stopped = false
	}
	if j.pgid != 0 {
		syscall.Kill(-j.pgid, syscall.SIGCONT)
	}
}

// builtinJobs lists the jobs.
func builtinJobs(stdout io.Writer) error {
	for _, j := range jobTable {
		j.wait(true)
	}
	for _, j := range append([]*job(nil), jobTable...) {
		fmt.Fprintln(stdout, j.describe(j.state()))
		if j.finished() {
			removeJob(j)
		}
	}
	return nil
}

// builtinFg continues a job in the foreground and waits for it.
func builtinFg(args []string, stdout io.Writer) error {
	if !jobControl {
		return errors.New("fg: no job control")
	}
	j, err := findJob("fg", args)
	if err != nil {
		return err
	}
	fmt.Fprintln(stdout, sanitize(j.text))

	if j.pgid != 0 {
		shellTTY, _ = readline.GetState(ttyFd)
		if j.stopped() && j.tty != nil {
			readline.Restore(ttyFd, j.tty)
		}
		setForeground(j.pgid)
	}
	continueJob(j)
	return waitForeground(j)
}

// builtinBg continues a stopped job in the background.
func builtinBg(args []string, stdout io.Writer) error {
	if !jobControl {
		return errors.New("bg: no job control")
	}
	j, err := findJob("bg", args)
	if err != nil {
		return err
	}
	if !j.stopped() {
		return fmt.Errorf("bg: job %d already in background", j.id)
	}
	continueJob(j)
	addJob(j) // the current job now
	fmt.Fprintln(stdout, j.describe("Running"))
	return nil
}

// builtinWait waits for a job, or for all background jobs. Stopped jobs are
// not waited for.
func builtinWait(args []string) error {
	jobs := append([]*job(nil), jobTable...)
	if len(args) > 1 {
		j, err := findJob("wait", args)
		if err != nil {
			return err
		}
		jobs = []*job{j}
	}

	var err error
	for _, j := range jobs {
		j.wait(false)
		err = nil
		if j.finished() {
			removeJob(j)
			err = j.result()
		}
	}
	return err
}

// exitShell exits the shell, unless there are stopped jobs and the user was
// not warned yet (like bash).
func exitShell(code int) error {
	if interactiveShell && hasStoppedJobs() && !exitWarned {
		exitWarned = true
		return errors.New("there are stopped jobs")
	}
	releaseTerminal()
	os.Exit(code)
	return nil
}