package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// The declared timeouts, which the session start sets (module.json's xinitrc contribution) and an
// inhibition returns to when it cannot read what was in force.
const (
	declaredLock    = 1800
	declaredStandby = 1800
	declaredSuspend = 1800
	declaredOff     = 3600

	mostInhibit = 600

	inhibitUnit = "screen-lock-inhibit"
	lockerUnit  = "screen-lock"
)

// clockTicks is the kernel's USER_HZ, which /proc/<pid>/stat counts a start time in: 100 on every
// architecture Arch Linux builds for.
const clockTicks = 100

func locker() string { return filepath.Join(operatorHome(), ".local", "bin", "screen-lock") }

// LockState is what lock and locked answer.
type LockState struct {
	Locked bool `json:"locked"`
	// Since is when the locker started, when it runs.
	Since      string `json:"since,omitempty"`
	PIDs       []int  `json:"pids"`
	LockedHint *bool  `json:"locked_hint,omitempty"`
	Watcher    bool   `json:"watcher_running"`
	Inhibited  bool   `json:"inhibited"`
	Via        string `json:"via,omitempty"`
}

// Lock locks through logind when the watcher runs, else runs the locker itself.
func Lock() (LockState, error) {
	s, err := findSession()
	if err != nil {
		return LockState{}, err
	}
	via := ""
	switch {
	case len(processesOf("i3lock")) > 0:
		via = "already locked"
	case len(processesOf("xss-lock")) > 0 && s.SessionID != "":
		r, err := s.run(10*time.Second, "", "loginctl", "lock-session", s.SessionID)
		if err != nil {
			return LockState{}, err
		}
		if r.Code != 0 {
			return LockState{}, fmt.Errorf("loginctl lock-session %s: %s", s.SessionID, strings.TrimSpace(r.Stderr))
		}
		via = "logind, answered by xss-lock"
	default:
		// No watcher: the session start's loop is not running (a session begun before this module
		// was assigned). The locker is run directly, under the account's service manager.
		if err := s.detach(lockerUnit, locker()); err != nil {
			return LockState{}, err
		}
		via = "the locker directly: xss-lock is not running in this session"
	}
	deadline := time.Now().Add(3 * time.Second)
	for len(processesOf("i3lock")) == 0 && time.Now().Before(deadline) {
		time.Sleep(100 * time.Millisecond)
	}
	state := lockState(s)
	state.Via = via
	if !state.Locked {
		return state, errors.New("asked to lock, and no locker is running 3s later")
	}
	return state, nil
}

// Locked answers the lock state without changing it.
func Locked() (LockState, error) {
	s := findEnvironment()
	return lockState(s), nil
}

func lockState(s Session) LockState {
	pids := processesOf("i3lock")
	st := LockState{Locked: len(pids) > 0, PIDs: pids, Watcher: len(processesOf("xss-lock")) > 0}
	if st.PIDs == nil {
		st.PIDs = []int{}
	}
	if len(pids) > 0 {
		if at, ok := startTime(pids[0]); ok {
			st.Since = at.Format(time.RFC3339)
		}
	}
	if s.SessionID != "" {
		if r, err := s.run(5*time.Second, "", "loginctl", "show-session", s.SessionID, "-p", "LockedHint", "--value"); err == nil && r.Code == 0 {
			hint := strings.TrimSpace(r.Stdout) == "yes"
			st.LockedHint = &hint
		}
	}
	if s.RuntimeDir != "" {
		if r, err := s.run(5*time.Second, "", "systemctl", "--user", "is-active", inhibitUnit+".service"); err == nil {
			st.Inhibited = strings.TrimSpace(r.Stdout) == "active"
		}
	}
	return st
}

// startTime is when a process started, from its start in clock ticks after boot and the boot time.
func startTime(pid int) (time.Time, bool) {
	stat, err := os.ReadFile(filepath.Join(procRoot, strconv.Itoa(pid), "stat"))
	if err != nil {
		return time.Time{}, false
	}
	// The command name is in parentheses and may hold spaces; the fields after it are fixed.
	end := strings.LastIndexByte(string(stat), ')')
	if end < 0 {
		return time.Time{}, false
	}
	fields := strings.Fields(string(stat[end+1:]))
	// starttime is field 22 of the whole line; after "pid (comm)" it is the 20th.
	if len(fields) < 20 {
		return time.Time{}, false
	}
	ticks, err := strconv.ParseInt(fields[19], 10, 64)
	if err != nil {
		return time.Time{}, false
	}
	boot, ok := bootTime()
	if !ok {
		return time.Time{}, false
	}
	return boot.Add(time.Duration(ticks) * time.Second / clockTicks), true
}

func bootTime() (time.Time, bool) {
	raw, err := os.ReadFile(filepath.Join(procRoot, "stat"))
	if err != nil {
		return time.Time{}, false
	}
	for _, line := range strings.Split(string(raw), "\n") {
		if v, ok := strings.CutPrefix(line, "btime "); ok {
			n, err := strconv.ParseInt(strings.TrimSpace(v), 10, 64)
			if err == nil {
				return time.Unix(n, 0), true
			}
		}
	}
	return time.Time{}, false
}

// IdleState is the screen's idle timeouts in force.
type IdleState struct {
	LockAfterSeconds int    `json:"lock_after_seconds"`
	CycleSeconds     int    `json:"cycle_seconds"`
	DPMSEnabled      bool   `json:"dpms_enabled"`
	StandbySeconds   int    `json:"standby_seconds"`
	SuspendSeconds   int    `json:"suspend_seconds"`
	OffSeconds       int    `json:"off_seconds"`
	MonitorOn        bool   `json:"monitor_on"`
	Declared         string `json:"declared"`
	Note             string `json:"note,omitempty"`
}

var (
	screensaverLine = regexp.MustCompile(`timeout:\s+(\d+)\s+cycle:\s+(\d+)`)
	dpmsLine        = regexp.MustCompile(`Standby:\s+(\d+)\s+Suspend:\s+(\d+)\s+Off:\s+(\d+)`)
)

func parseXsetQ(out string) (IdleState, error) {
	var st IdleState
	m := screensaverLine.FindStringSubmatch(out)
	if m == nil {
		return st, errors.New("xset q shows no screensaver timeout")
	}
	st.LockAfterSeconds, _ = strconv.Atoi(m[1])
	st.CycleSeconds, _ = strconv.Atoi(m[2])
	if d := dpmsLine.FindStringSubmatch(out); d != nil {
		st.StandbySeconds, _ = strconv.Atoi(d[1])
		st.SuspendSeconds, _ = strconv.Atoi(d[2])
		st.OffSeconds, _ = strconv.Atoi(d[3])
	}
	st.DPMSEnabled = strings.Contains(out, "DPMS is Enabled")
	st.MonitorOn = strings.Contains(out, "Monitor is On")
	st.Declared = fmt.Sprintf("lock after %ds; DPMS %d/%d/%d", declaredLock, declaredStandby, declaredSuspend, declaredOff)
	return st, nil
}

// IdleChange is what screen_lock_idle was asked to change; nil fields stay.
type IdleChange struct {
	LockAfter, Standby, Suspend, Off *int
}

func idleChangeOf(args map[string]any) (IdleChange, error) {
	var c IdleChange
	for key, into := range map[string]**int{
		"lock_after_seconds": &c.LockAfter, "standby_seconds": &c.Standby,
		"suspend_seconds": &c.Suspend, "off_seconds": &c.Off,
	} {
		if _, given := args[key]; !given {
			continue
		}
		n, err := whole(args, key, 0, 0, 24*3600)
		if err != nil {
			return c, err
		}
		*into = &n
	}
	return c, nil
}

func (c IdleChange) empty() bool {
	return c.LockAfter == nil && c.Standby == nil && c.Suspend == nil && c.Off == nil
}

// Idle reads the timeouts, and changes those asked for.
func Idle(change IdleChange) (IdleState, error) {
	s, err := findSession()
	if err != nil {
		return IdleState{}, err
	}
	before, err := xsetQ(s)
	if err != nil {
		return IdleState{}, err
	}
	if change.empty() {
		return before, nil
	}
	if change.LockAfter != nil {
		cycle := before.CycleSeconds
		if *change.LockAfter > 0 && cycle == 0 {
			cycle = *change.LockAfter
		}
		if err := xset(s, "s", strconv.Itoa(*change.LockAfter), strconv.Itoa(cycle)); err != nil {
			return IdleState{}, err
		}
	}
	if change.Standby != nil || change.Suspend != nil || change.Off != nil {
		pick := func(c *int, was int) string {
			if c != nil {
				return strconv.Itoa(*c)
			}
			return strconv.Itoa(was)
		}
		if err := xset(s, "dpms", pick(change.Standby, before.StandbySeconds), pick(change.Suspend, before.SuspendSeconds),
			pick(change.Off, before.OffSeconds)); err != nil {
			return IdleState{}, err
		}
	}
	after, err := xsetQ(s)
	if err != nil {
		return IdleState{}, err
	}
	after.Note = "changed for this session only; the declared timeouts return at the next login"
	return after, nil
}

func xsetQ(s Session) (IdleState, error) {
	r, err := s.run(5*time.Second, "", "xset", "q")
	if err != nil {
		return IdleState{}, err
	}
	if r.Code != 0 {
		return IdleState{}, fmt.Errorf("xset q: %s", strings.TrimSpace(r.Stderr))
	}
	return parseXsetQ(r.Stdout)
}

func xset(s Session, args ...string) error {
	r, err := s.run(5*time.Second, "", "xset", args...)
	if err != nil {
		return err
	}
	if r.Code != 0 {
		return fmt.Errorf("xset %s: %s", strings.Join(args, " "), strings.TrimSpace(r.Stderr))
	}
	return nil
}

// InhibitResult is what screen_lock_inhibit answers.
type InhibitResult struct {
	Inhibited bool   `json:"inhibited"`
	Until     string `json:"until,omitempty"`
	Restores  string `json:"restores,omitempty"`
}

// Inhibit keeps the screen on for minutes, then restores the timeouts that were in force; 0 ends an
// inhibition now. The waiting runs under the account's service manager, so it outlives this call,
// and stopping it restores at once.
func Inhibit(minutes int) (InhibitResult, error) {
	s, err := findSession()
	if err != nil {
		return InhibitResult{}, err
	}
	if minutes == 0 {
		r, err := s.run(10*time.Second, "", "systemctl", "--user", "stop", inhibitUnit+".service")
		if err != nil {
			return InhibitResult{}, err
		}
		if r.Code != 0 {
			return InhibitResult{}, fmt.Errorf("ending the inhibition: %s", strings.TrimSpace(r.Stderr))
		}
		return InhibitResult{Inhibited: false}, nil
	}
	// What to return to: what is in force now, unless an inhibition is already holding it at off.
	was, err := xsetQ(s)
	if err != nil {
		return InhibitResult{}, err
	}
	if lockState(s).Inhibited || (was.LockAfterSeconds == 0 && !was.DPMSEnabled) {
		was = IdleState{LockAfterSeconds: declaredLock, CycleSeconds: declaredLock, DPMSEnabled: true,
			StandbySeconds: declaredStandby, SuspendSeconds: declaredSuspend, OffSeconds: declaredOff}
	}
	script := inhibitScript(minutes, was)
	if err := s.detach(inhibitUnit, "/bin/sh", "-c", script); err != nil {
		return InhibitResult{}, err
	}
	return InhibitResult{Inhibited: true, Until: time.Now().Add(time.Duration(minutes) * time.Minute).Format(time.RFC3339),
		Restores: fmt.Sprintf("lock after %ds; DPMS %d/%d/%d", was.LockAfterSeconds, was.StandbySeconds, was.SuspendSeconds, was.OffSeconds)}, nil
}

func inhibitScript(minutes int, was IdleState) string {
	dpms := "+dpms"
	if !was.DPMSEnabled {
		dpms = "-dpms"
	}
	return fmt.Sprintf("restore() { xset s %d %d; xset dpms %d %d %d; xset %s; }; "+
		"trap 'restore; exit 0' TERM INT; xset s off -dpms; sleep %d & wait; restore",
		was.LockAfterSeconds, was.CycleSeconds, was.StandbySeconds, was.SuspendSeconds, was.OffSeconds, dpms, minutes*60)
}