// Runs cabletest for a fixed time and dumps the panel to png, since the machine
// with the framebuffer on it is not the machine looking at it.
//
//	sudo go run ./harness -for 20s -at 3s,15s -- -a enp1s0f0np0 -b enp1s0f1np1
package main

import (
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/unix"
)

const (
	fbPath   = "/dev/fb0"
	shotsDir = "shots"
	grace    = 3 * time.Second
)

// Matching fb.go, which reads the variable screen info as a flat run of u32s
// rather than a struct to sidestep any question of padding.
const (
	fbioGetVScreenInfo = 0x4600
	viXres             = 0
	viYres             = 1
	viBitsPerPixel     = 6
	viRedOffset        = 8
	viGreenOffset      = 11
	viBlueOffset       = 14
	viScreenInfoLen    = 40
)

func main() {
	runFor := flag.Duration("for", 15*time.Second, "how long to let cabletest run")
	at := flag.String("at", "", "offsets to capture the panel at, comma separated, e.g. 3s,10s")
	flag.Parse()

	if err := run(*runFor, *at, flag.Args()); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func parseShots(s string, runFor time.Duration) ([]time.Duration, error) {
	if s == "" {
		return nil, nil
	}
	var out []time.Duration
	for _, f := range strings.Split(s, ",") {
		d, err := time.ParseDuration(strings.TrimSpace(f))
		if err != nil {
			return nil, fmt.Errorf("bad offset %q: %w", f, err)
		}
		if d < 0 || d >= runFor {
			return nil, fmt.Errorf("offset %s falls outside the %s run", d, runFor)
		}
		out = append(out, d)
	}
	slices.Sort(out)
	return out, nil
}

func run(runFor time.Duration, at string, args []string) error {
	if os.Geteuid() != 0 {
		return fmt.Errorf("needs root for the raw sockets and the framebuffer: sudo go run ./harness")
	}
	shots, err := parseShots(at, runFor)
	if err != nil {
		return err
	}
	if len(shots) > 0 {
		if err := os.MkdirAll(shotsDir, 0o755); err != nil {
			return err
		}
	}

	cmd := exec.Command("go", append([]string{"run", "."}, args...)...)
	cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	// Its own process group: go run's compiled child is reparented rather than
	// killed when go run dies, and an orphan holding the wire poisons every
	// measurement after it.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		return err
	}
	pgid := cmd.Process.Pid

	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()

	start := time.Now()
	for _, d := range shots {
		select {
		case err := <-done:
			return exited(start, err)
		case <-time.After(time.Until(start.Add(d))):
		}
		name := filepath.Join(shotsDir, d.String()+".png")
		if err := shoot(name); err != nil {
			stop(pgid, done)
			return err
		}
		fmt.Printf("captured %s\n", name)
	}

	select {
	case err := <-done:
		return exited(start, err)
	case <-time.After(time.Until(start.Add(runFor))):
	}
	return stop(pgid, done)
}

func exited(start time.Time, err error) error {
	ran := time.Since(start).Round(time.Millisecond)
	if err != nil {
		return fmt.Errorf("cabletest exited after %s: %w", ran, err)
	}
	return fmt.Errorf("cabletest exited after %s", ran)
}

func stop(pgid int, done chan error) error {
	syscall.Kill(-pgid, syscall.SIGINT)
	select {
	case <-done:
	case <-time.After(grace):
		syscall.Kill(-pgid, syscall.SIGKILL)
		<-done
	}
	// Reaping go run says nothing about what it left behind, so the group itself
	// is what gets checked.
	if groupGone(pgid) {
		return nil
	}
	syscall.Kill(-pgid, syscall.SIGKILL)
	if groupGone(pgid) {
		return nil
	}
	return fmt.Errorf("process group %d survived SIGKILL; check with pgrep -a cabletest", pgid)
}

func groupGone(pgid int) bool {
	deadline := time.Now().Add(grace)
	for {
		if syscall.Kill(-pgid, 0) == syscall.ESRCH {
			return true
		}
		if time.Now().After(deadline) {
			return false
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func shoot(name string) error {
	f, err := os.Open(fbPath)
	if err != nil {
		return err
	}
	defer f.Close()

	var vi [viScreenInfoLen]uint32
	if _, _, errno := unix.Syscall(unix.SYS_IOCTL, f.Fd(),
		fbioGetVScreenInfo, uintptr(unsafe.Pointer(&vi[0]))); errno != 0 {
		return fmt.Errorf("get screen info: %w", errno)
	}
	if vi[viBitsPerPixel] != 32 {
		return fmt.Errorf("only 32bpp supported, got %d", vi[viBitsPerPixel])
	}
	stride, err := readUint("/sys/class/graphics/fb0/stride")
	if err != nil {
		return err
	}

	pw, ph := int(vi[viXres]), int(vi[viYres])
	mem, err := unix.Mmap(int(f.Fd()), 0, int(stride)*ph, unix.PROT_READ, unix.MAP_SHARED)
	if err != nil {
		return fmt.Errorf("mmap: %w", err)
	}
	defer unix.Munmap(mem)

	// cabletest flushes a whole frame into this mapping every 16ms, so the pixel
	// loop below is far too slow to read it directly: consecutive scanlines come
	// from different frames and every changing digit ends up drawn twice. One
	// memmove out first narrows the window to about the length of its own copy.
	buf := make([]byte, len(mem))
	copy(buf, mem)

	rs, gs, bs := vi[viRedOffset], vi[viGreenOffset], vi[viBlueOffset]
	// The panel is landscape and every draw is turned a quarter turn on its way
	// into it, so the turn is undone here to get back what a person standing in
	// front of it sees.
	img := image.NewNRGBA(image.Rect(0, 0, ph, pw))
	for x := 0; x < ph; x++ {
		for y := 0; y < pw; y++ {
			o := x*int(stride) + (pw-1-y)*4
			v := uint32(buf[o]) | uint32(buf[o+1])<<8 |
				uint32(buf[o+2])<<16 | uint32(buf[o+3])<<24
			img.SetNRGBA(x, y, color.NRGBA{
				R: uint8(v >> rs), G: uint8(v >> gs), B: uint8(v >> bs), A: 255,
			})
		}
	}

	out, err := os.Create(name)
	if err != nil {
		return err
	}
	if err := png.Encode(out, img); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

func readUint(path string) (uint64, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	v, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64)
	if err != nil {
		return 0, fmt.Errorf("%s: %w", path, err)
	}
	return v, nil
}