// Runs cabletest for a fixed time and dumps the panel to png, since the machine
// with the display 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 (
	"bytes"
	"flag"
	"fmt"
	"image/png"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"syscall"
	"time"
)

const (
	shotsDir = "shots"
	grace    = 3 * time.Second
	// The framebuffer's DRM master can outlive the process group by a moment;
	// an immediate next run finds the device busy.
	fuzzGap = 2 * time.Second
)

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")
	fuzzN := flag.Int("fuzz", 0, "repeat the whole lifecycle this many times; full output only for runs that fail or log module events")
	flag.Parse()

	var err error
	if *fuzzN > 0 {
		err = fuzz(*fuzzN, *runFor, flag.Args())
	} else {
		err = run(*runFor, *at, os.Stdout, flag.Args())
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func fuzz(n int, runFor time.Duration, args []string) error {
	failed := 0
	for i := 1; i <= n; i++ {
		if i > 1 {
			time.Sleep(fuzzGap)
		}
		var buf bytes.Buffer
		start := time.Now()
		err := run(runFor, "", &buf, args)
		dur := time.Since(start).Round(100 * time.Millisecond)
		if err != nil {
			failed++
			name := fmt.Sprintf("fuzz-run%02d.log", i)
			if werr := os.WriteFile(name, buf.Bytes(), 0o644); werr != nil {
				name = fmt.Sprintf("unsaved: %v", werr)
			}
			fmt.Printf("run %02d/%02d: FAILED (%s): %v [%s]\n", i, n, dur, err, name)
			os.Stdout.Write(buf.Bytes())
			continue
		}
		fmt.Printf("run %02d/%02d: ok (%s)\n", i, n, dur)
	}
	fmt.Printf("fuzz: %d/%d ok, %d failed\n", n-failed, n, failed)
	if failed > 0 {
		return fmt.Errorf("%d/%d runs failed", failed, n)
	}
	return nil
}

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, out io.Writer, 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 = out, out
	// 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 {
	g, err := openGrabber()
	if err != nil {
		return err
	}
	defer g.close()

	img, err := g.frame()
	if err != nil {
		return err
	}

	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()
}