// stayup: a nohup replacement backed by systemd.
//
// Every command becomes a real systemd service unit (Restart=always), so it
// survives crashes, OOM kills and host reboots. stayup keeps no state of its
// own: the unit files under systemd's unit directory are the only record.
//
//   - Non-root: ~/.config/systemd/user/stayup-<cmd>-<hash>.service via
//     `systemctl --user`; `loginctl enable-linger` makes it start at boot.
//   - Root: /etc/systemd/system/, WantedBy=multi-user.target.
//   - The unit captures cwd, the resolved executable and the caller's
//     environment (file mode 0600). Logs go to the journal unless -o/--stderr
//     are given (StandardOutput=append:, systemd >= 240).
//   - Same command + same cwd = same unit, so starting twice is idempotent.
//   - Users see "processes"; unit names only surface as NAME in -l.
package main

import (
	"bufio"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"text/tabwriter"
	"time"
	"unsafe"
)

const prefix = "stayup-"

var (
	user    = os.Geteuid() != 0 // non-root: per-user systemd instance
	unitDir string
)

func main() {
	var (
		fg      bool
		attach  bool
		list    bool
		kill    bool
		outFile string
		errFile string
	)
	flag.BoolVar(&fg, "f", false, "")
	flag.BoolVar(&fg, "foreground", false, "follow the log; ctrl-c kills, ctrl-d detaches")
	flag.BoolVar(&attach, "a", false, "")
	flag.BoolVar(&attach, "attach", false, "attach to a running process (by command, PID or name)")
	flag.BoolVar(&list, "l", false, "")
	flag.BoolVar(&list, "list", false, "list processes")
	flag.BoolVar(&kill, "K", false, "")
	flag.BoolVar(&kill, "kill", false, "kill a process (by command, PID or name)")
	flag.StringVar(&outFile, "o", "", "")
	flag.StringVar(&outFile, "stdout", "", "append stdout to this file (default: system log)")
	flag.StringVar(&errFile, "stderr", "", "append stderr to this file (default: same as stdout)")
	flag.Usage = usage
	flag.Parse()
	args := flag.Args()

	if user {
		cfg := os.Getenv("XDG_CONFIG_HOME")
		if cfg == "" {
			cfg = filepath.Join(os.Getenv("HOME"), ".config")
		}
		unitDir = filepath.Join(cfg, "systemd", "user")
	} else {
		unitDir = "/etc/systemd/system"
	}

	switch {
	case list:
		fail(doList())
	case kill:
		if len(args) == 0 {
			fail(fmt.Errorf("--kill needs a command or PID"))
		}
		u, err := find(args)
		fail(err)
		c := command(u)
		fail(remove(u))
		fmt.Fprintln(os.Stderr, "killed", c)
	case attach:
		if len(args) == 0 {
			fail(fmt.Errorf("--attach needs a command or PID"))
		}
		u, err := find(args)
		fail(err)
		fail(foreground(u))
	case len(args) == 0:
		usage()
		os.Exit(2)
	default:
		if errFile == "" {
			errFile = outFile
		}
		u, fresh, err := start(args, outFile, errFile)
		fail(err)
		if !fg {
			if fresh {
				fmt.Fprintf(os.Stderr, "started %s (pid %s)\n", shJoin(args), show(u, "MainPID"))
			} else {
				fmt.Fprintf(os.Stderr, "already running: %s (pid %s)\n", shJoin(args), show(u, "MainPID"))
			}
			fmt.Fprintf(os.Stderr, "attach: stayup -a %s\nkill:   stayup -K %s\n", shJoin(args), shJoin(args))
			return
		}
		fail(foreground(u))
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `stayup - keep a command running via systemd (survives crashes, OOM, reboots)

  stayup [-f] [-o FILE] [--stderr FILE] [--] COMMAND [ARGS...]
  stayup -a | --attach COMMAND [ARGS...] | PID | NAME
  stayup -l | --list
  stayup -K | --kill COMMAND [ARGS...] | PID | NAME

  -f, --foreground   run attached: ctrl-c kills, ctrl-d or hangup detaches
  -a, --attach       attach to a running process (same keys as -f)
  -o, --stdout FILE  append stdout to FILE (default: system log)
      --stderr FILE  append stderr to FILE (default: same as stdout)
  -l, --list         list processes
  -K, --kill         kill a process (by command, PID or NAME from -l)
`)
}

func fail(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, "stayup:", err)
		os.Exit(1)
	}
}

// ---- systemctl plumbing ---------------------------------------------------

func sysctl(args ...string) *exec.Cmd {
	if user {
		args = append([]string{"--user"}, args...)
	}
	c := exec.Command("systemctl", args...)
	c.Stderr = os.Stderr
	return c
}

func run(args ...string) error {
	if err := sysctl(args...).Run(); err != nil {
		return fmt.Errorf("systemctl %s: %w", strings.Join(args, " "), err)
	}
	return nil
}

// command returns the user's command line for a unit.
func command(unit string) string {
	return strings.TrimPrefix(show(unit, "Description"), "stayup: ")
}

// show returns one property of a unit, or "" when it is unknown.
func show(unit, prop string) string {
	out, _ := sysctl("show", "-p", prop, "--value", unit).Output()
	return strings.TrimSpace(string(out))
}

func props(unit string, names ...string) map[string]string {
	m := map[string]string{}
	out, _ := sysctl("show", "-p", strings.Join(names, ","), unit).Output()
	for _, l := range strings.Split(string(out), "\n") {
		if k, v, ok := strings.Cut(l, "="); ok {
			m[k] = v
		}
	}
	return m
}

// ---- unit naming -----------------------------------------------------------

var badChars = regexp.MustCompile(`[^a-zA-Z0-9_.]+`)

func unitName(cwd string, args []string) string {
	h := sha256.Sum256([]byte(cwd + "\x00" + strings.Join(args, "\x00")))
	slug := badChars.ReplaceAllString(filepath.Base(args[0]), "-")
	if len(slug) > 24 {
		slug = slug[:24]
	}
	return prefix + slug + "-" + hex.EncodeToString(h[:])[:8] + ".service"
}

func shJoin(args []string) string {
	q := make([]string, len(args))
	for i, a := range args {
		if a == "" || strings.ContainsAny(a, " \t\n'\"\\$&|;<>()*?[]#~") {
			a = "'" + strings.ReplaceAll(a, "'", `'\''`) + "'"
		}
		q[i] = a
	}
	return strings.Join(q, " ")
}

// unitQuote quotes one word for ExecStart=/Environment= in a unit file.
func unitQuote(s string) string {
	r := strings.NewReplacer(`\`, `\\`, `"`, `\"`, "%", "%%", "\n", `\n`)
	return `"` + r.Replace(s) + `"`
}

func allUnits() ([]string, error) {
	ents, err := os.ReadDir(unitDir)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var us []string
	for _, e := range ents {
		if n := e.Name(); strings.HasPrefix(n, prefix) && strings.HasSuffix(n, ".service") {
			us = append(us, n)
		}
	}
	return us, nil
}

// find resolves a --kill/--attach target: a unit name, a PID, the same command
// in this cwd, or the same command in any cwd when unambiguous.
func find(args []string) (string, error) {
	units, err := allUnits()
	if err != nil {
		return "", err
	}
	if len(args) == 1 {
		name := args[0]
		if !strings.HasSuffix(name, ".service") {
			name += ".service"
		}
		for _, u := range units {
			if u == name {
				return u, nil
			}
		}
		if pid, err := strconv.Atoi(args[0]); err == nil {
			for _, u := range units {
				if show(u, "MainPID") == strconv.Itoa(pid) {
					return u, nil
				}
			}
			return "", fmt.Errorf("no process with pid %d (see stayup -l)", pid)
		}
	}
	cwd, _ := os.Getwd()
	want := unitName(cwd, args)
	var matches []string
	for _, u := range units {
		if u == want {
			return u, nil
		}
		if show(u, "Description") == "stayup: "+shJoin(args) {
			matches = append(matches, u)
		}
	}
	switch len(matches) {
	case 0:
		return "", fmt.Errorf("no process for %q (see stayup -l)", shJoin(args))
	case 1:
		return matches[0], nil
	}
	return "", fmt.Errorf("%d processes run %q in other directories; use the PID or NAME from stayup -l", len(matches), shJoin(args))
}

// ---- actions ---------------------------------------------------------------

func start(args []string, outFile, errFile string) (unit string, fresh bool, err error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", false, err
	}
	unit = unitName(cwd, args)
	path := filepath.Join(unitDir, unit)
	if _, err := os.Stat(path); err == nil {
		if show(unit, "ActiveState") != "active" {
			err = run("start", unit)
		}
		return unit, false, err
	}

	exe, err := exec.LookPath(args[0])
	if err != nil {
		return "", false, err
	}
	if exe, err = filepath.Abs(exe); err != nil {
		return "", false, err
	}
	abs := func(f string) string {
		if f == "" || filepath.IsAbs(f) {
			return f
		}
		return filepath.Join(cwd, f)
	}
	outFile, errFile = abs(outFile), abs(errFile)

	var b strings.Builder
	fmt.Fprintf(&b, "# generated by stayup on %s\n", time.Now().Format(time.RFC3339))
	fmt.Fprintf(&b, "[Unit]\nDescription=stayup: %s\nStartLimitIntervalSec=0\n\n", shJoin(args))
	fmt.Fprintf(&b, "[Service]\nType=simple\nWorkingDirectory=%s\n", strings.ReplaceAll(cwd, "%", "%%"))
	words := []string{unitQuote(exe)}
	for _, a := range args[1:] {
		words = append(words, unitQuote(a))
	}
	fmt.Fprintf(&b, "ExecStart=%s\n", strings.Join(words, " "))
	fmt.Fprint(&b, "Restart=always\nRestartSec=1\nTimeoutStopSec=15\nOOMPolicy=continue\n")
	if outFile != "" {
		fmt.Fprintf(&b, "StandardOutput=append:%s\n", outFile)
	}
	if errFile != "" {
		fmt.Fprintf(&b, "StandardError=append:%s\n", errFile)
	}
	for _, kv := range os.Environ() {
		if k, _, _ := strings.Cut(kv, "="); k != "" && k != "_" && !strings.HasPrefix(k, "BASH_FUNC_") {
			fmt.Fprintf(&b, "Environment=%s\n", unitQuote(kv))
		}
	}
	target := "multi-user.target"
	if user {
		target = "default.target"
	}
	fmt.Fprintf(&b, "\n[Install]\nWantedBy=%s\n", target)

	if err := os.MkdirAll(unitDir, 0o755); err != nil {
		return "", false, err
	}
	if err := os.WriteFile(path, []byte(b.String()), 0o600); err != nil {
		return "", false, err
	}
	if user {
		// Let the user's systemd instance (and our service) outlive logins and come back after reboot.
		exec.Command("loginctl", "enable-linger").Run()
	}
	if err := run("daemon-reload"); err != nil {
		return "", false, err
	}
	if err := run("enable", "--now", "--quiet", unit); err != nil {
		sysctl("disable", unit).Run()
		os.Remove(path)
		run("daemon-reload")
		return "", false, fmt.Errorf("%w\nunit file was:\n%s", err, b.String())
	}
	return unit, true, nil
}

func remove(unit string) error {
	err := run("disable", "--now", "--quiet", unit)
	q := sysctl("reset-failed", unit)
	q.Stderr = nil
	q.Run()
	if rerr := os.Remove(filepath.Join(unitDir, unit)); rerr != nil && !os.IsNotExist(rerr) && err == nil {
		err = rerr
	}
	run("daemon-reload")
	return err
}

func doList() error {
	units, err := allUnits()
	if err != nil {
		return err
	}
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
	fmt.Fprintln(w, "NAME\tPID\tSTATE\tSTARTED\tLAST RESTART\tRESTARTS\tDIR\tCOMMAND")
	for _, u := range units {
		p := props(u, "MainPID", "ActiveState", "SubState", "NRestarts", "ExecMainStartTimestamp", "Description", "WorkingDirectory")
		started := "?"
		if fi, err := os.Stat(filepath.Join(unitDir, u)); err == nil {
			started = fi.ModTime().Format("2006-01-02 15:04")
		}
		pid := p["MainPID"]
		if pid == "0" || pid == "" {
			pid = "-"
		}
		fmt.Fprintf(w, "%s\t%s\t%s/%s\t%s\t%s\t%s\t%s\t%s\n",
			strings.TrimSuffix(u, ".service"), pid, p["ActiveState"], p["SubState"], started, fmtTS(p["ExecMainStartTimestamp"]),
			p["NRestarts"], p["WorkingDirectory"], strings.TrimPrefix(p["Description"], "stayup: "))
	}
	return w.Flush()
}

// fmtTS trims systemd's "Day YYYY-MM-DD HH:MM:SS TZ" to "YYYY-MM-DD HH:MM".
func fmtTS(s string) string {
	if f := strings.Fields(s); len(f) >= 3 {
		return f[1] + " " + f[2][:5]
	}
	return "-"
}

// ---- foreground ------------------------------------------------------------

const hint = "ctrl-c to kill, ctrl-d or close connection to hang up and stay up"

func foreground(unit string) error {
	// Log follower: journalctl, or tail -F for file outputs.
	var follow *exec.Cmd
	p := props(unit, "StandardOutput", "StandardError")
	var files []string
	for _, v := range []string{p["StandardOutput"], p["StandardError"]} {
		if f, ok := strings.CutPrefix(v, "append:"); ok && (len(files) == 0 || files[0] != f) {
			files = append(files, f)
		}
	}
	if len(files) > 0 {
		follow = exec.Command("tail", append([]string{"-n", "50", "-F"}, files...)...)
	} else {
		ja := []string{"-u", unit, "-f", "-n", "50", "-o", "cat"}
		if user {
			ja = append([]string{"--user"}, ja...)
		}
		follow = exec.Command("journalctl", ja...)
	}
	follow.Stdout, follow.Stderr = os.Stdout, os.Stderr
	follow.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} // keep terminal ctrl-c for us
	if err := follow.Start(); err != nil {
		return err
	}
	defer follow.Process.Kill()

	tty := isTerminal(os.Stderr.Fd())
	if tty {
		pinHint(true)
		defer unpinHint()
	} else {
		fmt.Fprintln(os.Stderr, hint)
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGWINCH)
	eof := make(chan struct{})
	if isTerminal(os.Stdin.Fd()) {
		go func() {
			r := bufio.NewReader(os.Stdin)
			for {
				if _, err := r.ReadByte(); err != nil {
					close(eof)
					return
				}
			}
		}()
	}
	done := make(chan error, 1)
	go func() { done <- follow.Wait() }()

	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGWINCH:
				if tty {
					pinHint(false)
				}
				continue
			case syscall.SIGINT:
				unpinHint()
				fmt.Fprintln(os.Stderr, "\nkilling", command(unit))
				return remove(unit)
			default: // SIGHUP, SIGTERM: detach, leave it running
				return nil
			}
		case <-eof:
			unpinHint()
			fmt.Fprintf(os.Stderr, "\ndetached; %s stays up\n", command(unit))
			return nil
		case err := <-done:
			return fmt.Errorf("log output ended: %v", err)
		}
	}
}

// Terminal helpers: keep the hint on the last row using a scroll region.

func isTerminal(fd uintptr) bool {
	var t syscall.Termios
	_, _, e := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(&t)))
	return e == 0
}

func rows() int {
	var ws struct{ Row, Col, X, Y uint16 }
	syscall.Syscall(syscall.SYS_IOCTL, os.Stderr.Fd(), syscall.TIOCGWINSZ, uintptr(unsafe.Pointer(&ws)))
	if ws.Row == 0 {
		return 24
	}
	return int(ws.Row)
}

func pinHint(clear bool) {
	h := rows()
	if clear {
		fmt.Fprint(os.Stderr, "\x1b[2J")
	}
	// scroll region = rows 1..h-1, hint on row h (reverse video), cursor back into the region
	fmt.Fprintf(os.Stderr, "\x1b[1;%dr\x1b[%d;1H\x1b[2K\x1b[7m%s\x1b[0m\x1b[%d;1H", h-1, h, hint, h-1)
}

func unpinHint() {
	if isTerminal(os.Stderr.Fd()) {
		fmt.Fprintf(os.Stderr, "\x1b[r\x1b[%d;1H\x1b[2K", rows())
	}
}