// Runs cabletest until interrupted or for -for, dumps the panel to png, and
// streams it as png deltas on :8080, 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"
	"errors"
	"flag"
	"fmt"
	"image/png"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"slices"
	"strings"
	"syscall"
	"time"

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

var errInterrupted = errors.New("interrupted")

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", 0, "how long to let cabletest run; 0 runs until interrupted")
	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()

	if err := serveStream(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	restoreTerm := rawKeys()
	go relayStdin()

	var err error
	if *fuzzN > 0 {
		if *runFor <= 0 {
			restoreTerm()
			fmt.Fprintln(os.Stderr, "fuzz needs a bounded -for")
			os.Exit(1)
		}
		err = fuzz(*fuzzN, *runFor, flag.Args())
	} else {
		err = run(*runFor, *at, os.Stdout, flag.Args())
	}
	restoreTerm()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// Console keys reach cabletest as they are pressed, the same as at the
// appliance; isig stays on so ^C still stops the run through the relay.
func rawKeys() func() {
	fd := int(os.Stdin.Fd())
	orig, err := unix.IoctlGetTermios(fd, unix.TCGETS)
	if err != nil {
		return func() {}
	}
	saved := *orig
	raw := *orig
	raw.Lflag &^= unix.ICANON | unix.ECHO
	raw.Cc[unix.VMIN] = 1
	raw.Cc[unix.VTIME] = 0
	if unix.IoctlSetTermios(fd, unix.TCSETS, &raw) != nil {
		return func() {}
	}
	return func() { unix.IoctlSetTermios(fd, unix.TCSETS, &saved) }
}

// The child's own process group makes the tty off limits to it (reading from
// the background raises SIGTTIN), so the harness reads and forwards to the
// run that currently owns the console. One reader for the life of the
// process: a copier per run would pile up under fuzz, each stealing
// keystrokes for a run already gone.
func relayStdin() {
	buf := make([]byte, 64)
	for {
		n, err := os.Stdin.Read(buf)
		if err != nil {
			return
		}
		consoleMu.Lock()
		if console != nil {
			console.Write(buf[:n])
		}
		consoleMu.Unlock()
	}
}

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 errors.Is(err, errInterrupted) {
			os.Stdout.Write(buf.Bytes())
			return err
		}
		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 || (runFor > 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
		}
	}

	// The child gets its own process group below, which is exactly what stops
	// the terminal's ^C reaching it, so the harness has to relay the stop.
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	defer signal.Stop(sigs)

	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}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	consoleMu.Lock()
	console = stdin
	consoleMu.Unlock()
	defer func() {
		consoleMu.Lock()
		console = nil
		consoleMu.Unlock()
	}()
	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 <-sigs:
			return interrupted(start, pgid, done)
		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)
	}

	var expire <-chan time.Time
	if runFor > 0 {
		expire = time.After(time.Until(start.Add(runFor)))
	}
	select {
	case err := <-done:
		return exited(start, err)
	case <-sigs:
		return interrupted(start, pgid, done)
	case <-expire:
	}
	return stop(pgid, done)
}

func interrupted(start time.Time, pgid int, done chan error) error {
	if err := stop(pgid, done); err != nil {
		return err
	}
	return fmt.Errorf("%w after %s", errInterrupted, time.Since(start).Round(time.Millisecond))
}

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