// Package testsuite provides many quick-and-dirty integration testing utilities.
//
// Attempting to import this package outside testing causes the resulting
// program to panic.
package testsuite

import (
	"bufio"
	"context"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"hakurei.app/check"
	"hakurei.app/fhs"
	"hakurei.app/hst"
	"hakurei.app/internal/store"
)

// serial is the next value returned by [NewSerial].
var serial atomic.Uint64

// NewSerial returns a serial number.
func NewSerial() uint64 { defer serial.Add(1); return serial.Load() }

// ReceiveSignals blocks until a termination signal arrives, and terminates.
func ReceiveSignals() {
	s := make(chan os.Signal, 3)
	signal.Notify(s, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
	log.Fatalf("terminating on signal %s", <-s)
}

var (
	// users caches [user.LookupId] calls.
	users = make(map[string]*user.User)
	// usersMu synchronises access to users.
	usersMu sync.RWMutex
)

// mustLookupId is like [user.LookupId], but the first result is cached.
func mustLookupId(uid string) *user.User {
	usersMu.RLock()
	v, ok := users[uid]
	usersMu.RUnlock()
	if ok {
		return v
	}

	u, err := user.LookupId(uid)
	if err != nil {
		log.Fatal(err)
	}

	usersMu.Lock()
	users[uid] = u
	usersMu.Unlock()
	return u
}

// MustAppendEnv adds extra environment variables to cmd.
func MustAppendEnv(cmd *exec.Cmd, env ...string) {
	if len(cmd.Env) == 0 {
		var cred syscall.Credential
		if cmd.SysProcAttr != nil && cmd.SysProcAttr.Credential != nil {
			cred = *cmd.SysProcAttr.Credential
		}
		u := mustLookupId(strconv.Itoa(int(cred.Uid)))

		cmd.Env = append(cmd.Env,
			"PATH="+os.Getenv("PATH"),
			"HOME="+u.HomeDir,
			"USER="+u.Username,
			"USERNAME="+u.Username,
		)
	}
	cmd.Env = append(cmd.Env, env...)
}

// MustRun runs command and terminates the testsuite on error.
func MustRun(cred *syscall.Credential, extraEnv []string, command ...string) {
	cmd := exec.Command(command[0], command[1:]...)
	cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig:  syscall.SIGKILL,
		Credential: cred,
	}
	MustAppendEnv(cmd, extraEnv...)
	if err := cmd.Run(); err != nil {
		log.Fatalf("must run %s: %v", strings.Join(command, " "), err)
	}
}

// ErrUnexpectedSuccess is returned for processes expected to exit with a
// non-zero code, but failed to do so.
var ErrUnexpectedSuccess = errors.New("process unexpectedly exited with code 0")

// MustFail runs command and terminates the testsuite if the program fails to
// start or exits with code 0.
func MustFail(
	cred *syscall.Credential,
	extraEnv []string,
	quiet bool,
	command ...string,
) {
	cmd := exec.Command(command[0], command[1:]...)
	if !quiet {
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	}
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig:  syscall.SIGKILL,
		Credential: cred,
	}
	MustAppendEnv(cmd, extraEnv...)
	if err := cmd.Run(); err == nil {
		log.Fatal(ErrUnexpectedSuccess)
	} else if e, ok := errors.AsType[*exec.ExitError](err); !ok {
		log.Fatalf("must fail %s: %v", strings.Join(command, " "), err)
	} else if !e.Exited() {
		log.Fatalf("must fail %s: %v", strings.Join(command, " "), e)
	}
}

// MustStart starts cmd and returns a channel delivering its wait error.
func MustStart(cmd *exec.Cmd) (done <-chan error) {
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	d := make(chan error, 1)
	go func() { d <- cmd.Wait() }()
	return d
}

// MustStartWith wraps [MustStart] and creates the [exec.Cmd] object internally.
func MustStartWith(
	ctx context.Context,
	cred *syscall.Credential,
	quiet bool,
	extraEnv []string,
	files []*os.File,
	command ...string,
) (proc *os.Process, done <-chan error) {
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	if !quiet {
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	}
	cmd.ExtraFiles = files
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig:  syscall.SIGTERM,
		Credential: cred,
	}
	MustAppendEnv(cmd, extraEnv...)
	return cmd.Process, MustStart(cmd)
}

// MustOutput is like [MustRun] and [MustFail], but collects the entire output
// in memory and returns them as strings.
func MustOutput(
	ctx context.Context,
	cred *syscall.Credential,
	extraEnv []string,
	exitCode int,
	name string,
	command ...string,
) (string, string) {
	var stdout, stderr strings.Builder
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig:  syscall.SIGTERM,
		Credential: cred,
	}
	MustAppendEnv(cmd, extraEnv...)

	err := cmd.Run()
	if err == nil {
		if exitCode == 0 {
			return stdout.String(), stderr.String()
		}
		log.Fatalf(
			"%s succeeded unexpectedly"+
				"\nstdout = %q\nstderr = %q",
			name,
			stdout.String(),
			stderr.String(),
		)
	}

	if e, ok := errors.AsType[*exec.ExitError](err); !ok {
		log.Fatal(err)
	} else if e.ExitCode() != exitCode {
		log.Fatalf(
			"%s: %v"+
				"\nstdout = %q\nstderr = %q",
			name, e,
			stdout.String(),
			stderr.String(),
		)
	}
	return stdout.String(), stderr.String()
}

// AssertOutput returns a function that terminates the program if called with
// arguments differing from the ones it is initialised with.
func AssertOutput(action, stdout, stderr string) (
	assert func(stdout, stderr string),
) {
	return func(gotStdout, gotStderr string) {
		if stdout != gotStdout || stderr != gotStderr {
			log.Fatalf(
				"unexpected output %s"+
					"\nstdout = %q\nstderr = %q",
				action,
				gotStdout,
				gotStderr,
			)
		}
	}
}

// MustCheckFilter is like [CheckFilter], but terminates the test suite if a
// non-nil error is returned. Otherwise, the tracee is terminated after it
// resumes.
func MustCheckFilter(pid int, sum [sha512.Size]byte) {
	// podman installs its own filter
	if err := CheckFilter(pid, 1, sum); err != nil {
		log.Fatal(err)
	} else if err = syscall.Kill(pid, syscall.SIGTERM); err != nil {
		log.Fatalf("cannot terminate tracee: %v", err)
	}
}

// FilterTerminated returns a non-nil error if err is not an [exec.ExitError]
// describing a process terminated by a syscall.SIGTERM signal.
func FilterTerminated(err error) error {
	if err == nil {
		return ErrUnexpectedSuccess
	}

	e, ok := errors.AsType[*exec.ExitError](err)
	if !ok {
		return err
	}

	if e.ExitCode() == 0x80+int(syscall.SIGTERM) {
		return nil
	}
	return e
}

// Poll repeatedly runs command until it succeeds.
func Poll(
	d time.Duration,
	cred *syscall.Credential,
	extraEnv []string,
	command ...string,
) {
	for range time.NewTicker(d).C {
		cmd := exec.Command(command[0], command[1:]...)
		cmd.SysProcAttr = &syscall.SysProcAttr{
			Pdeathsig:  syscall.SIGKILL,
			Credential: cred,
		}
		MustAppendEnv(cmd, extraEnv...)
		if err := cmd.Run(); err != nil {
			if e, ok := errors.AsType[*exec.ExitError](err); ok && e.Exited() {
				continue
			}
			log.Fatalf("poll %s: %v", strings.Join(command, " "), err)
		}
		break
	}
}

const (
	// XDGRuntimeDir is the hardcoded XDG runtime directory of the only user
	// mapped in hsurc.
	XDGRuntimeDir = "/var/run/user/1000"

	// XDGRuntimeEnv is the environment variable string for XDG_RUNTIME_DIR.
	XDGRuntimeEnv = "XDG_RUNTIME_DIR=" + XDGRuntimeDir

	// HakureiPrefix is the hakurei work directory of the only user mapped in
	// hsurc.
	HakureiPrefix = "/tmp/hakurei.0/"
)

// MustPrepareUserDirs creates user-specific directories.
func MustPrepareUserDirs(cred *syscall.Credential) {
	if err := os.MkdirAll(XDGRuntimeDir, 0700); err != nil {
		log.Fatal(err)
	} else if err = os.Chown(
		XDGRuntimeDir,
		int(cred.Uid),
		int(cred.Gid),
	); err != nil {
		log.Fatal(err)
	}
}

// SessionEnv returns environment variables referring to the hosted session.
func SessionEnv(dbusEnv string) []string {
	return []string{
		"PATH=" + os.Getenv("PATH"),
		"TERM=xterm",
		XDGRuntimeEnv,
		WaylandEnv,
		"DISPLAY=:0",
		dbusEnv,
	}
}

// quietServices is initialised from the environment and controls whether
// session services should have their outputs connected.
var quietServices = os.Getenv("HAKUREI_VERBOSE_SERVICES") == ""

// MustStartSessionBus starts a session bus that is never explicitly terminated.
// The test suite is terminated if the session bus daemon terminates.
func MustStartSessionBus(cred *syscall.Credential) (dbusEnv string) {
	r, w, err := os.Pipe()
	if err != nil {
		log.Fatal(err)
	}

	// this is never explicitly terminated
	_, done := MustStartWith(
		context.Background(), cred, quietServices, nil, []*os.File{w},
		"dbus-daemon",
		"--print-address=3",
		"--address=unix:path="+XDGRuntimeDir+"/dbus",
		"--session",
		"--nofork",
		"--nopidfile",
	)

	go func() {
		if _err := <-done; _err != nil {
			log.Fatalf("session bus terminated unexpectedly: %v", _err)
		}
		log.Fatal("session bus terminated unexpectedly")
	}()

	dbusEnv, err = bufio.NewReader(r).ReadString('\n')
	if err != nil {
		log.Fatal(err)
	}
	dbusEnv = dbusEnv[:len(dbusEnv)-1]
	log.Printf("dbus listening on %s", dbusEnv)
	dbusEnv = "DBUS_SESSION_BUS_ADDRESS=" + dbusEnv

	if err = r.Close(); err != nil {
		log.Fatal(err)
	}
	return
}

const (
	// SwayEnv is the environment variable string for the sway IPC socket.
	SwayEnv = "SWAYSOCK=" + XDGRuntimeDir + "/sway"
	// WaylandEnv is the environment variable string for the wayland display.
	WaylandEnv = "WAYLAND_DISPLAY=wayland-1"
)

// MustStartSway starts the sway wayland display server which must be terminated
// by calling [TerminateSway].
func MustStartSway(
	wg *sync.WaitGroup,
	cred *syscall.Credential,
	dbusEnv string,
) {
	wg.Go(func() {
		// this is terminated via swaymsg
		_, done := MustStartWith(
			context.Background(), cred, quietServices, []string{
				"WLR_BACKENDS=headless",
				XDGRuntimeEnv,
				SwayEnv,
				dbusEnv,
			}, nil,
			"sway",
		)
		if err := <-done; err != nil {
			log.Fatalf("sway terminated unexpectedly: %v", err)
		}
	})

	Poll(
		50*time.Millisecond,
		cred,
		[]string{SwayEnv},
		"swaymsg",
	)
	log.Printf("sway available via %s", SwayEnv)
}

// TerminateSway requests for the sway server to terminate via sway IPC.
func TerminateSway(cred *syscall.Credential) {
	MustFail(cred, []string{SwayEnv}, quietServices, "swaymsg", "exit")
}

// MustStartPipeWire starts a PipeWire server that is never explicitly
// terminated.  The test suite is terminated if the PipeWire server terminates.
func MustStartPipeWire(cred *syscall.Credential, dbusEnv string) {
	// this is never explicitly terminated
	_, done := MustStartWith(
		context.Background(), cred, quietServices, []string{
			XDGRuntimeEnv,
			dbusEnv,
		}, nil,
		"pipewire",
	)

	go func() {
		if _err := <-done; _err != nil {
			log.Fatalf("pipewire terminated unexpectedly: %v", _err)
		}
		log.Fatal("pipewire terminated unexpectedly")
	}()

	Poll(50*time.Millisecond, cred, []string{
		XDGRuntimeEnv,
		dbusEnv,
	},
		"wpctl",
		"status",
	)

	_, _done := MustStartWith(
		context.Background(), cred, quietServices, []string{
			XDGRuntimeEnv,
			dbusEnv,
		}, nil,
		"wireplumber",
	)

	go func() {
		if _err := <-_done; _err != nil {
			log.Fatalf("wireplumber terminated unexpectedly: %v", _err)
		}
		log.Fatal("wireplumber terminated unexpectedly")
	}()
}

// MustScanFor continuously scans the proc filesystem and calls f for each entry
// visited.
func MustScanFor(f func(ps *StatScanner) bool) int {
	var ps StatScanner

	for ps.Scan() {
		if f(&ps) {
			break
		}
	}
	if err := ps.Err(); err != nil {
		log.Fatal(err)
	}
	return ps.Stat().PID
}

// StartHakurei starts a hakurei container configured by c.
func StartHakurei(
	ctx context.Context,
	wg *sync.WaitGroup,
	cred *syscall.Credential,
	dbusEnv string,
	processCommand func(cmd *exec.Cmd),
	onDone func(err error, output string),
	scan, wait bool,
	c *hst.Config,
) (p hst.State, pid int, err error) {
	s := store.New(check.MustAbs("/tmp/hakurei.0/state"))

	cmd := exec.CommandContext(ctx,
		"script", "/dev/null",
		"-E", "always",
		"-qec",
		"hakurei run "+
			"--identifier-fd=5"+
			" 4 1>&3",
	)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig:  syscall.SIGTERM,
		Credential: cred,
	}
	var output strings.Builder
	cmd.Stdout, cmd.Stderr = &output, &output
	cmd.Env = SessionEnv(dbusEnv)

	var notify, _notify, _conf, conf, ident, _ident *os.File
	if notify, _notify, err = os.Pipe(); err != nil {
		return
	}
	cmd.ExtraFiles = append(cmd.ExtraFiles, _notify)
	if _conf, conf, err = os.Pipe(); err != nil {
		_, _ = notify.Close(), _notify.Close()
		return
	}
	cmd.ExtraFiles = append(cmd.ExtraFiles, _conf)
	if ident, _ident, err = os.Pipe(); err != nil {
		_, _ = notify.Close(), _notify.Close()
		_, _ = conf.Close(), _conf.Close()
		return
	}
	cmd.ExtraFiles = append(cmd.ExtraFiles, _ident)
	if processCommand != nil {
		processCommand(cmd)
	}

	done := MustStart(cmd)
	wg.Go(func() { _err := <-done; onDone(_err, output.String()) })

	if err = errors.Join(
		_notify.Close(),
		_conf.Close(),
		_ident.Close(),
	); err != nil {
		_, _, _ = notify.Close(), conf.Close(), ident.Close()
		return
	}

	if err = json.NewEncoder(conf).Encode(c); err != nil {
		_, _, _ = notify.Close(), conf.Close(), ident.Close()
		return
	} else if err = conf.Close(); err != nil {
		_, _ = notify.Close(), ident.Close()
		return
	}

	if !scan && len(cmd.ExtraFiles) != 3 {
		err = errors.Join(notify.Close(), ident.Close())
		return
	}

	var id hst.ID
	if _, err = io.ReadFull(ident, id[:]); err != nil {
		_, _ = notify.Close(), ident.Close()
		return
	} else if err = ident.Close(); err != nil {
		_ = notify.Close()
		return
	}

	if !scan {
		err = notify.Close()
		return
	}

	if wait {
		if _, err = io.ReadFull(notify, make([]byte, 8)); err != nil {
			_ = notify.Close()
			return
		}
	}
	if err = notify.Close(); err != nil {
		return
	}

	var ok bool
	entries, copyError := s.All()
	for entry := range entries {
		if entry.ID == id {
			ok = true
			if _, err = entry.Load(&p, nil); err != nil {
				log.Fatal(err)
			}
			break
		}
	}
	if err = copyError(); err != nil {
		log.Fatal(err)
	}
	if !ok {
		log.Fatalf("instance %s is not present in store", id)
	}

	var stat syscall.Stat_t
	pid = MustScanFor(func(ps *StatScanner) bool {
		select {
		case err = <-done:
			if err == nil {
				log.Fatal("test process terminated unexpectedly")
			}
			log.Fatal(err)
		default:
			break
		}

		var args []string
		if args, err = ps.Stat().Args(); err != nil {
			if IsNotExist(err) {
				return false
			}
			log.Fatal(err)
		} else if !slices.Equal(args, c.Container.Args) {
			return false
		}

		if err = ps.Stat().Stat(&stat); err != nil {
			if IsNotExist(err) {
				return false
			}
			log.Fatal(err)
		}

		uid := hst.ToUser[uint32](0, uint32(c.Identity))
		if stat.Uid != uid || stat.Gid != uid {
			return false
		}

		var t []byte
		if t, err = os.ReadFile(filepath.Join(
			fhs.Proc,
			strconv.Itoa(ps.Stat().PPID),
			"stat",
		)); err != nil {
			if IsNotExist(err) {
				return false
			}
			log.Fatal(err)
		}

		var _stat Stat
		if err = _stat.UnmarshalText(t); err != nil {
			log.Fatal(err)
		}
		if _stat.PPID != p.ShimPID {
			return false
		}

		return true
	})
	return
}

// MustStartHakurei is like [StartHakurei], but terminates the test suite if
// an error is encountered.
func MustStartHakurei(
	ctx context.Context,
	wg *sync.WaitGroup,
	cred *syscall.Credential,
	dbusEnv string,
	onDone func(err error, output string),
	wait bool,
	c *hst.Config,
) (p hst.State, pid int) {
	_ctx, cancel := context.WithCancel(ctx)

	var err error
	p, pid, err = StartHakurei(
		_ctx, wg, cred, dbusEnv,
		nil, func(err error, output string) {
			defer cancel()
			onDone(err, output)
		}, true, wait, c,
	)
	if err != nil {
		log.Printf("cannot start hakurei container: %v", err)
		cancel()
	}
	return
}