// The operator's graphical session, as a tool the node's runtime runs finds it (novox/hq ADR 0208).
//
// The runtime is a system service running as the operator account (ADR 0175): it has the account's
// uid and none of the session's environment — no DISPLAY, no XAUTHORITY, no session bus. A tool that
// draws on the screen or talks to the desktop's D-Bus must find them. It reads them from a process of
// the account that is part of the session (the window manager first), the same thing `loginctl` and
// a person's own shell would point at, and says where it found them.
//
// Long-lived programs a tool starts go to the account's own service manager through `systemd-run
// --user`, never as children of the tool: the runtime's unit is a cgroup the service manager empties
// whenever the runtime restarts, and a compositor or a clipboard owner started from inside it would
// die with it.
//
// This file is the same in every desktop module that carries it; it moves into the Go SDK once a
// second consumer outside the desktop wants it.
package main

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Where the session is looked for. Variables so a test can point them at a fake tree.
var (
	procRoot   = "/proc"
	runUserDir = "/run/user"
	x11Sockets = "/tmp/.X11-unix"
)

// sessionHolders are the processes whose environment is the session's, best first: the window
// manager is the session, the rest are its children. Anything else carrying DISPLAY ranks after them.
var sessionHolders = []string{"i3", "sway", "i3bar", "picom", "xss-lock", "dunst", "clipmenud", "xterm"}

// sessionKeys are the variables a session carries that a tool hands on to what it runs.
var sessionKeys = []string{"DISPLAY", "XAUTHORITY", "WAYLAND_DISPLAY", "DBUS_SESSION_BUS_ADDRESS",
	"XDG_RUNTIME_DIR", "XDG_SESSION_ID", "I3SOCK"}

// Session is what a tool needs to reach the operator's desktop.
type Session struct {
	UID        int    `json:"uid"`
	Display    string `json:"display,omitempty"`
	XAuthority string `json:"xauthority,omitempty"`
	Wayland    string `json:"wayland_display,omitempty"`
	Bus        string `json:"bus,omitempty"`
	RuntimeDir string `json:"runtime_dir,omitempty"`
	SessionID  string `json:"session_id,omitempty"`
	I3Sock     string `json:"i3sock,omitempty"`
	// From says where the values were found: the tool's own environment, a process, or the socket.
	From string `json:"from"`
}

// ErrNoSession is answered by a tool that needs the desktop when nobody is logged in to it.
var ErrNoSession = errors.New("no graphical session")

// ErrTimedOut is what run answers for a command ended because it ran past its time.
var ErrTimedOut = errors.New("timed out")

// ErrNoBus is answered by a tool that needs the session bus when the account has none.
var ErrNoBus = errors.New("no session bus")

// operatorHome is the account's home: what the runtime was told, else the process's own.
func operatorHome() string {
	if h := strings.TrimSpace(os.Getenv("MESH_OPERATOR_HOME")); h != "" {
		return h
	}
	h, _ := os.UserHomeDir()
	return h
}

// findSession finds the graphical session of the account this tool runs as, or answers
// ErrNoSession with what it looked at.
func findSession() (Session, error) {
	s := findEnvironment()
	if s.Display == "" && s.Wayland == "" {
		return s, fmt.Errorf("%w for uid %d on this machine: no process of the account carries DISPLAY "+
			"or WAYLAND_DISPLAY, and no X server socket in %s has an authority file to go with it. "+
			"Is anyone logged in to the desktop?", ErrNoSession, s.UID, x11Sockets)
	}
	return s, nil
}

// findBus finds the account's session bus, which a logged-in account has whether or not a desktop
// is running.
func findBus() (Session, error) {
	s := findEnvironment()
	if s.Bus == "" {
		return s, fmt.Errorf("%w for uid %d: DBUS_SESSION_BUS_ADDRESS is not set and %s does not exist "+
			"(the account is not logged in)", ErrNoBus, s.UID, filepath.Join(runUserDir, strconv.Itoa(s.UID), "bus"))
	}
	return s, nil
}

func findEnvironment() Session {
	uid := os.Getuid()
	s := Session{UID: uid}
	own := map[string]string{}
	for _, k := range sessionKeys {
		own[k] = os.Getenv(k)
	}
	if own["DISPLAY"] != "" || own["WAYLAND_DISPLAY"] != "" {
		s.fill(own)
		s.From = "the tool's own environment"
	} else if pid, comm, env, ok := sessionProcess(uid); ok {
		s.fill(env)
		s.From = fmt.Sprintf("process %s (pid %d)", comm, pid)
	} else if display, ok := lonelyX11Socket(); ok {
		if a := filepath.Join(operatorHome(), ".Xauthority"); exists(a) {
			s.Display, s.XAuthority = display, a
			s.From = "the X server socket and the account's ~/.Xauthority"
		}
		s.fill(own)
	} else {
		s.fill(own)
		s.From = "nothing: no session found"
	}
	// The bus and the runtime directory are the account's, whether or not the process named them.
	runtime := filepath.Join(runUserDir, strconv.Itoa(uid))
	if s.RuntimeDir == "" && exists(runtime) {
		s.RuntimeDir = runtime
	}
	if s.Bus == "" && s.RuntimeDir != "" && exists(filepath.Join(s.RuntimeDir, "bus")) {
		s.Bus = "unix:path=" + filepath.Join(s.RuntimeDir, "bus")
	}
	return s
}

func (s *Session) fill(env map[string]string) {
	set := func(dst *string, key string) {
		if *dst == "" {
			*dst = env[key]
		}
	}
	set(&s.Display, "DISPLAY")
	set(&s.XAuthority, "XAUTHORITY")
	set(&s.Wayland, "WAYLAND_DISPLAY")
	set(&s.Bus, "DBUS_SESSION_BUS_ADDRESS")
	set(&s.RuntimeDir, "XDG_RUNTIME_DIR")
	set(&s.SessionID, "XDG_SESSION_ID")
	set(&s.I3Sock, "I3SOCK")
}

// sessionProcess is the best process of this uid whose environment names a display.
func sessionProcess(uid int) (int, string, map[string]string, bool) {
	entries, err := os.ReadDir(procRoot)
	if err != nil {
		return 0, "", nil, false
	}
	type candidate struct {
		pid  int
		comm string
		env  map[string]string
		rank int
	}
	var found []candidate
	for _, e := range entries {
		pid, err := strconv.Atoi(e.Name())
		if err != nil {
			continue
		}
		dir := filepath.Join(procRoot, e.Name())
		if owner, ok := ownerOf(dir); !ok || owner != uid {
			continue
		}
		raw, err := os.ReadFile(filepath.Join(dir, "environ"))
		if err != nil {
			continue
		}
		env := parseEnviron(raw)
		if env["DISPLAY"] == "" && env["WAYLAND_DISPLAY"] == "" {
			continue
		}
		comm := readTrimmed(filepath.Join(dir, "comm"))
		rank := len(sessionHolders)
		for i, h := range sessionHolders {
			if h == comm {
				rank = i
				break
			}
		}
		found = append(found, candidate{pid, comm, env, rank})
	}
	if len(found) == 0 {
		return 0, "", nil, false
	}
	sort.Slice(found, func(i, j int) bool {
		if found[i].rank != found[j].rank {
			return found[i].rank < found[j].rank
		}
		return found[i].pid > found[j].pid // the newer of two equals
	})
	best := found[0]
	return best.pid, best.comm, best.env, true
}

func parseEnviron(raw []byte) map[string]string {
	env := map[string]string{}
	for _, kv := range bytes.Split(raw, []byte{0}) {
		if i := bytes.IndexByte(kv, '='); i > 0 {
			env[string(kv[:i])] = string(kv[i+1:])
		}
	}
	return env
}

func ownerOf(path string) (int, bool) {
	info, err := os.Stat(path)
	if err != nil {
		return 0, false
	}
	st, ok := info.Sys().(*syscall.Stat_t)
	if !ok {
		return 0, false
	}
	return int(st.Uid), true
}

// lonelyX11Socket is the display of the one X server socket there is, when there is exactly one.
func lonelyX11Socket() (string, bool) {
	entries, err := os.ReadDir(x11Sockets)
	if err != nil {
		return "", false
	}
	var displays []string
	for _, e := range entries {
		if n := strings.TrimPrefix(e.Name(), "X"); n != e.Name() {
			if _, err := strconv.Atoi(n); err == nil {
				displays = append(displays, ":"+n)
			}
		}
	}
	if len(displays) != 1 {
		return "", false
	}
	return displays[0], true
}

func readTrimmed(path string) string {
	b, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

func exists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// Env is this process's environment with the session's variables in place of its own.
func (s Session) Env() []string {
	drop := map[string]bool{}
	for _, k := range sessionKeys {
		drop[k] = true
	}
	var env []string
	for _, kv := range os.Environ() {
		if i := strings.IndexByte(kv, '='); i > 0 && drop[kv[:i]] {
			continue
		}
		env = append(env, kv)
	}
	add := func(k, v string) {
		if v != "" {
			env = append(env, k+"="+v)
		}
	}
	add("DISPLAY", s.Display)
	add("XAUTHORITY", s.XAuthority)
	add("WAYLAND_DISPLAY", s.Wayland)
	add("DBUS_SESSION_BUS_ADDRESS", s.Bus)
	add("XDG_RUNTIME_DIR", s.RuntimeDir)
	add("XDG_SESSION_ID", s.SessionID)
	add("I3SOCK", s.I3Sock)
	return env
}

// mostOutput bounds what a command may answer with, per stream.
const mostOutput = 256 << 10

// Result is what a command did.
type Result struct {
	Stdout    string `json:"stdout"`
	Stderr    string `json:"stderr,omitempty"`
	Code      int    `json:"code"`
	Truncated bool   `json:"truncated,omitempty"`
}

// run runs a command in the session's environment, its input given, ended with everything it
// started after timeout. A command that is not installed is an error naming it; one that exits
// non-zero is a Result with its code, for the caller to judge.
func (s Session) run(timeout time.Duration, stdin string, name string, args ...string) (Result, error) {
	path, err := exec.LookPath(name)
	if err != nil {
		return Result{}, fmt.Errorf("%s is not installed on this machine", name)
	}
	cmd := exec.Command(path, args...)
	cmd.Env = s.Env()
	if home := operatorHome(); exists(home) {
		cmd.Dir = home
	}
	if stdin != "" {
		cmd.Stdin = strings.NewReader(stdin)
	}
	var out, errOut capped
	cmd.Stdout, cmd.Stderr = &out, &errOut
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		return Result{}, fmt.Errorf("%s: %w", name, err)
	}
	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()
	select {
	case err = <-done:
	case <-time.After(timeout):
		_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		<-done
		return Result{Stdout: out.String(), Stderr: errOut.String()},
			fmt.Errorf("%s did not finish within %s and was ended: %w", name, timeout, ErrTimedOut)
	}
	r := Result{Stdout: out.String(), Stderr: errOut.String(), Truncated: out.cut || errOut.cut}
	var exit *exec.ExitError
	if errors.As(err, &exit) {
		r.Code = exit.ExitCode()
	} else if err != nil {
		return r, fmt.Errorf("%s: %w", name, err)
	}
	return r, nil
}

// detach starts a long-lived program under the account's own service manager, as a transient unit
// that carries the session's display, so it outlives the runtime that asked for it. A unit already
// running under the same name is stopped first, so a fixed name means "at most one".
func (s Session) detach(unit string, args ...string) error {
	if s.RuntimeDir == "" {
		return fmt.Errorf("%w: the account's runtime directory is missing, so its service manager "+
			"cannot be reached", ErrNoBus)
	}
	_, _ = s.run(5*time.Second, "", "systemctl", "--user", "stop", unit+".service")
	call := []string{"--user", "--collect", "--quiet", "--unit=" + unit}
	for _, kv := range [][2]string{{"DISPLAY", s.Display}, {"XAUTHORITY", s.XAuthority},
		{"WAYLAND_DISPLAY", s.Wayland}, {"XDG_SESSION_ID", s.SessionID}, {"I3SOCK", s.I3Sock}} {
		if kv[1] != "" {
			call = append(call, "--setenv="+kv[0]+"="+kv[1])
		}
	}
	call = append(call, "--")
	call = append(call, args...)
	r, err := s.run(10*time.Second, "", "systemd-run", call...)
	if err != nil {
		return err
	}
	if r.Code != 0 {
		return fmt.Errorf("systemd-run %s: %s", unit, strings.TrimSpace(r.Stderr))
	}
	return nil
}

// uniqueUnit is a transient unit name that will not collide with an earlier one.
func uniqueUnit(prefix string) string {
	return fmt.Sprintf("%s-%d", prefix, time.Now().UnixNano())
}

type capped struct {
	bytes.Buffer
	cut bool
}

func (c *capped) Write(p []byte) (int, error) {
	if room := mostOutput - c.Len(); room < len(p) {
		if room > 0 {
			c.Buffer.Write(p[:room])
		}
		c.cut = true
		return len(p), nil
	}
	return c.Buffer.Write(p)
}

// processesOf are the pids of this uid's processes whose command name is comm, oldest first.
func processesOf(comm string) []int {
	entries, err := os.ReadDir(procRoot)
	if err != nil {
		return nil
	}
	uid := os.Getuid()
	var pids []int
	for _, e := range entries {
		pid, err := strconv.Atoi(e.Name())
		if err != nil {
			continue
		}
		dir := filepath.Join(procRoot, e.Name())
		if owner, ok := ownerOf(dir); !ok || owner != uid {
			continue
		}
		if readTrimmed(filepath.Join(dir, "comm")) == comm {
			pids = append(pids, pid)
		}
	}
	sort.Ints(pids)
	return pids
}

// signalAll sends sig to every process of this uid named comm, and answers the pids it reached.
func signalAll(comm string, sig syscall.Signal) []int {
	var reached []int
	for _, pid := range processesOf(comm) {
		if syscall.Kill(pid, sig) == nil {
			reached = append(reached, pid)
		}
	}
	return reached
}