// Package desktop is how a desktop module's tools act in the operator's graphical session
// (novox/hq ADR 0208, research 026/05).
//
// **One question, answered once for every desktop tool.** A tool runs inside the node's runtime: a
// process of node-tools.service, started by the system's service manager as the operator account,
// with no session around it — no DISPLAY, no XAUTHORITY, no session bus. The session it must act in
// was started elsewhere, by the login manager, and the only place its values are written down is
// the environment of the processes it started. So this package finds the session the way a person
// would: it looks at the operator account's own processes, takes the one that is plainly the
// session's (the window manager, or the oldest process carrying a display), confirms with logind
// that its session is a live local one, and checks that the display's socket is really there.
//
// **Only the session's own words are read.** A session's processes also carry whatever its start
// script exported — on the workstations that was a file of secrets — so the environment is filtered
// to a fixed list of names while it is read, and nothing else ever leaves /proc.
//
// The D-Bus address handed on is the user manager's socket, `unix:path=$XDG_RUNTIME_DIR/bus`,
// whenever it exists, because that is where the portal, the notifier and every user service
// listen. A session started on a private bus (a stale session, measured on one workstation) is
// reported as `session_bus` beside it, so the difference is visible rather than guessed at.
//
// The same copy of this package is vendored into every desktop module (xorg, lemurs, i3, xterm,
// adwaita); the catalogue builds each module alone, so it cannot be imported across them. Change
// every copy together — the modules' tests compare them.
package desktop

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

// SessionWords are the only environment words read from a session's process: the ones that say
// where the session is. Everything else in that environment is the operator's, and is never read.
var SessionWords = []string{
	"DISPLAY", "WAYLAND_DISPLAY", "XAUTHORITY",
	"XDG_SESSION_ID", "XDG_SESSION_TYPE", "XDG_SESSION_DESKTOP", "XDG_CURRENT_DESKTOP",
	"XDG_RUNTIME_DIR", "DBUS_SESSION_BUS_ADDRESS", "XDG_SEAT", "XDG_VTNR",
	"I3SOCK", "SWAYSOCK",
}

// Session is the operator's running graphical session, as a tool needs it.
type Session struct {
	UID            int    `json:"uid"`
	ID             string `json:"session_id,omitempty"`
	Type           string `json:"type"`
	Display        string `json:"display,omitempty"`
	WaylandDisplay string `json:"wayland_display,omitempty"`
	XAuthority     string `json:"xauthority,omitempty"`
	RuntimeDir     string `json:"runtime_dir"`
	Bus            string `json:"bus,omitempty"`
	SessionBus     string `json:"session_bus,omitempty"`
	Desktop        string `json:"desktop,omitempty"`
	// FoundIn is the process whose environment named the session.
	FoundIn Process `json:"found_in"`
	// Active is logind's word on the session, when logind answered.
	Active *bool `json:"active,omitempty"`

	words map[string]string
}

// Process is one process the search looked at.
type Process struct {
	PID     int    `json:"pid"`
	Command string `json:"command"`
	start   uint64
}

// NoSession is the answer when there is no graphical session to act in. Its text is JSON, so a tool
// that returns it as its error still answers structured data.
type NoSession struct {
	Reason string   `json:"reason"`
	Looked []string `json:"looked"`
}

func (e *NoSession) Error() string {
	b, _ := json.Marshal(map[string]any{"error": "no-graphical-session", "reason": e.Reason, "looked": e.Looked})
	return string(b)
}

// IsNoSession is whether err says there is no session.
func IsNoSession(err error) bool {
	var n *NoSession
	return errors.As(err, &n)
}

// Finder holds where the search looks, so a test can point it at a tree of its own.
type Finder struct {
	Proc        string // the process table: /proc
	X11Sockets  string // where X servers listen: /tmp/.X11-unix
	RuntimeBase string // the parent of every XDG_RUNTIME_DIR: /run/user
	UID         int    // whose session
	// Prefer names the processes that are the session's own, best first: the session's holder.
	Prefer []string
	// Logind answers `loginctl show-session` for one id; nil skips the check.
	Logind func(id string) (map[string]string, error)
}

// DefaultFinder is the machine's: the account this process runs as, or — when it runs as root — the
// operator account the runtime names (MESH_OPERATOR_ACCOUNT).
func DefaultFinder(prefer ...string) Finder {
	uid := os.Getuid()
	if uid == 0 {
		if name := os.Getenv("MESH_OPERATOR_ACCOUNT"); name != "" {
			if u, err := user.Lookup(name); err == nil {
				if n, err := strconv.Atoi(u.Uid); err == nil {
					uid = n
				}
			}
		}
	}
	return Finder{
		Proc: "/proc", X11Sockets: "/tmp/.X11-unix", RuntimeBase: "/run/user",
		UID: uid, Prefer: prefer, Logind: loginctl,
	}
}

// Find is the operator's session on this machine, preferring a process named in prefer.
func Find(prefer ...string) (*Session, error) {
	return DefaultFinder(prefer...).Find()
}

type candidate struct {
	proc   Process
	words  map[string]string
	rank   int
	logind map[string]string
}

// Find looks for the session.
func (f Finder) Find() (*Session, error) {
	entries, err := os.ReadDir(f.Proc)
	if err != nil {
		return nil, &NoSession{Reason: "the process table cannot be read: " + err.Error(), Looked: []string{f.Proc}}
	}
	looked := []string{fmt.Sprintf("the processes of uid %d in %s", f.UID, f.Proc)}
	var found []candidate
	stale := 0
	for _, e := range entries {
		pid, err := strconv.Atoi(e.Name())
		if err != nil {
			continue
		}
		dir := filepath.Join(f.Proc, e.Name())
		info, err := os.Stat(dir)
		if err != nil {
			continue
		}
		if st, ok := info.Sys().(*syscall.Stat_t); !ok || int(st.Uid) != f.UID {
			continue
		}
		words := readWords(filepath.Join(dir, "environ"))
		if words["DISPLAY"] == "" && words["WAYLAND_DISPLAY"] == "" {
			continue
		}
		if !f.reachable(words) {
			stale++
			continue
		}
		found = append(found, candidate{proc: Process{PID: pid, Command: comm(dir), start: startTime(dir)}, words: words})
	}
	if len(found) == 0 {
		reason := fmt.Sprintf("no process of uid %d carries a display", f.UID)
		if stale > 0 {
			reason = fmt.Sprintf("%d process(es) of uid %d name a display whose socket is gone: the session they belonged to has ended", stale, f.UID)
		}
		return nil, &NoSession{Reason: reason, Looked: append(looked, f.X11Sockets, f.RuntimeBase)}
	}

	// logind's word on each session the candidates name, asked once per session.
	asked := map[string]map[string]string{}
	for i := range found {
		id := found[i].words["XDG_SESSION_ID"]
		if f.Logind == nil || id == "" {
			found[i].rank = 1
			continue
		}
		props, done := asked[id]
		if !done {
			props, _ = f.Logind(id)
			asked[id] = props
		}
		found[i].logind = props
		switch {
		case props == nil:
			found[i].rank = 1
		case props["Remote"] == "yes":
			found[i].rank = 3
		case props["Active"] == "yes" && props["State"] != "closing":
			found[i].rank = 0
		case props["State"] == "closing":
			found[i].rank = 3
		default:
			found[i].rank = 2
		}
	}
	if f.Logind != nil {
		looked = append(looked, "logind's sessions")
	}
	preferred := func(c candidate) int {
		for i, p := range f.Prefer {
			if c.proc.Command == p {
				return i
			}
		}
		return len(f.Prefer)
	}
	sort.SliceStable(found, func(i, j int) bool {
		a, b := found[i], found[j]
		if a.rank != b.rank {
			return a.rank < b.rank
		}
		if pa, pb := preferred(a), preferred(b); pa != pb {
			return pa < pb
		}
		if a.proc.start != b.proc.start {
			return a.proc.start < b.proc.start
		}
		return a.proc.PID < b.proc.PID
	})
	best := found[0]
	if best.rank == 3 {
		return nil, &NoSession{Reason: "the only sessions found are remote or closing", Looked: looked}
	}
	return f.session(best), nil
}

func (f Finder) session(c candidate) *Session {
	w := c.words
	s := &Session{
		UID: f.UID, ID: w["XDG_SESSION_ID"], Display: w["DISPLAY"], WaylandDisplay: w["WAYLAND_DISPLAY"],
		XAuthority: w["XAUTHORITY"], RuntimeDir: w["XDG_RUNTIME_DIR"], FoundIn: c.proc, words: w,
	}
	s.Desktop = w["XDG_CURRENT_DESKTOP"]
	if s.Desktop == "" {
		s.Desktop = w["XDG_SESSION_DESKTOP"]
	}
	switch {
	case w["XDG_SESSION_TYPE"] != "":
		s.Type = w["XDG_SESSION_TYPE"]
	case s.WaylandDisplay != "":
		s.Type = "wayland"
	default:
		s.Type = "x11"
	}
	if s.RuntimeDir == "" {
		s.RuntimeDir = filepath.Join(f.RuntimeBase, strconv.Itoa(f.UID))
	}
	if isSocket(filepath.Join(s.RuntimeDir, "bus")) {
		s.Bus = "unix:path=" + filepath.Join(s.RuntimeDir, "bus")
	}
	if own := w["DBUS_SESSION_BUS_ADDRESS"]; own != "" && own != s.Bus {
		s.SessionBus = own
	}
	if c.logind != nil {
		active := c.logind["Active"] == "yes"
		s.Active = &active
	}
	return s
}

// reachable is whether the display a process names is still served: the X server's socket, or the
// Wayland compositor's. A process outliving its session still carries the session's words.
func (f Finder) reachable(w map[string]string) bool {
	if d := w["WAYLAND_DISPLAY"]; d != "" {
		path := d
		if !filepath.IsAbs(d) {
			dir := w["XDG_RUNTIME_DIR"]
			if dir == "" {
				dir = filepath.Join(f.RuntimeBase, strconv.Itoa(f.UID))
			}
			path = filepath.Join(dir, d)
		}
		if isSocket(path) {
			return true
		}
	}
	n, ok := DisplayNumber(w["DISPLAY"])
	return ok && isSocket(filepath.Join(f.X11Sockets, "X"+strconv.Itoa(n)))
}

// DisplayNumber is the server number of a local X display (":1", ":1.0", "unix:1"); a display on
// another host — an ssh session's forwarded one — is not the local session and answers false.
func DisplayNumber(display string) (int, bool) {
	host, rest, ok := strings.Cut(display, ":")
	if !ok || (host != "" && host != "unix") {
		return 0, false
	}
	num, _, _ := strings.Cut(rest, ".")
	n, err := strconv.Atoi(num)
	if err != nil || n < 0 {
		return 0, false
	}
	return n, true
}

// Word is one of the session's words as its process had it ("" when it had none).
func (s *Session) Word(name string) string { return s.words[name] }

// Env is base with the session's words in place of whatever base said for them.
func (s *Session) Env(base []string) []string {
	drop := map[string]bool{}
	for _, w := range SessionWords {
		drop[w] = true
	}
	out := make([]string, 0, len(base)+8)
	for _, kv := range base {
		k, _, _ := strings.Cut(kv, "=")
		if !drop[k] {
			out = append(out, kv)
		}
	}
	bus := s.Bus
	if bus == "" {
		bus = s.SessionBus
	}
	for _, kv := range [][2]string{
		{"DISPLAY", s.Display}, {"WAYLAND_DISPLAY", s.WaylandDisplay}, {"XAUTHORITY", s.XAuthority},
		{"XDG_RUNTIME_DIR", s.RuntimeDir}, {"DBUS_SESSION_BUS_ADDRESS", bus},
		{"XDG_SESSION_TYPE", s.Type}, {"XDG_SESSION_ID", s.ID},
		{"XDG_CURRENT_DESKTOP", s.words["XDG_CURRENT_DESKTOP"]},
		{"XDG_SESSION_DESKTOP", s.words["XDG_SESSION_DESKTOP"]},
		{"I3SOCK", s.words["I3SOCK"]}, {"SWAYSOCK", s.words["SWAYSOCK"]},
	} {
		if kv[1] != "" {
			out = append(out, kv[0]+"="+kv[1])
		}
	}
	return out
}

// UserEnv is base with the account's own runtime directory and bus, for a tool that talks to the
// user manager or the session bus and needs no display — it works with no session at all.
func UserEnv(base []string, uid int) []string {
	dir := filepath.Join("/run/user", strconv.Itoa(uid))
	out := make([]string, 0, len(base)+2)
	for _, kv := range base {
		k, _, _ := strings.Cut(kv, "=")
		if k != "XDG_RUNTIME_DIR" && k != "DBUS_SESSION_BUS_ADDRESS" {
			out = append(out, kv)
		}
	}
	return append(out, "XDG_RUNTIME_DIR="+dir, "DBUS_SESSION_BUS_ADDRESS=unix:path="+filepath.Join(dir, "bus"))
}

// readWords reads a process's environment and keeps only SessionWords.
func readWords(path string) map[string]string {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	keep := map[string]bool{}
	for _, w := range SessionWords {
		keep[w] = true
	}
	out := map[string]string{}
	for _, kv := range bytes.Split(raw, []byte{0}) {
		k, v, ok := bytes.Cut(kv, []byte{'='})
		if ok && keep[string(k)] {
			out[string(k)] = string(v)
		}
	}
	return out
}

func comm(dir string) string {
	b, err := os.ReadFile(filepath.Join(dir, "comm"))
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// startTime is field 22 of /proc/<pid>/stat: when the process started, in clock ticks since boot.
// Read after the command's closing parenthesis, because the command may hold spaces.
func startTime(dir string) uint64 {
	b, err := os.ReadFile(filepath.Join(dir, "stat"))
	if err != nil {
		return ^uint64(0)
	}
	i := bytes.LastIndexByte(b, ')')
	if i < 0 {
		return ^uint64(0)
	}
	fields := strings.Fields(string(b[i+1:]))
	// fields[0] is the state, field 3 of the line; start time is field 22.
	if len(fields) < 20 {
		return ^uint64(0)
	}
	n, err := strconv.ParseUint(fields[19], 10, 64)
	if err != nil {
		return ^uint64(0)
	}
	return n
}

func isSocket(path string) bool {
	info, err := os.Stat(path)
	return err == nil && info.Mode()&os.ModeSocket != 0
}

// loginctl asks logind about one session, by its property lines.
func loginctl(id string) (map[string]string, error) {
	cmd := exec.Command("loginctl", "show-session", id, "-p", "Active", "-p", "State", "-p", "Remote", "-p", "Type", "-p", "Class")
	var out bytes.Buffer
	cmd.Stdout = &out
	done := make(chan error, 1)
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	go func() { done <- cmd.Wait() }()
	select {
	case err := <-done:
		if err != nil {
			return nil, err
		}
	case <-time.After(3 * time.Second):
		_ = cmd.Process.Kill()
		return nil, errors.New("loginctl did not answer in 3s")
	}
	return ParseProperties(out.String()), nil
}

// ParseProperties reads `Key=Value` lines, as loginctl and systemctl show print them.
func ParseProperties(text string) map[string]string {
	out := map[string]string{}
	sc := bufio.NewScanner(strings.NewReader(text))
	for sc.Scan() {
		if k, v, ok := strings.Cut(sc.Text(), "="); ok {
			out[k] = v
		}
	}
	return out
}