picomap/cmd/picomap/main.go

package main

import (
	"encoding/hex"
	"flag"
	"fmt"
	"log/slog"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"github.com/theater/picomap/lib/client"
	"github.com/theater/picomap/lib/picotool"
)

func main() {
	if len(os.Args) < 2 {
		slog.Error("usage: picomap <info|load|test> [args...]")
		os.Exit(1)
	}
	cmd := os.Args[1]
	args := os.Args[2:]

	var err error
	switch cmd {
	case "info":
		err = cmdInfo(args)
	case "load":
		err = cmdLoad(args)
	case "log":
		err = cmdLog(args)
	case "test":
		err = cmdTestGroup(args)
	default:
		slog.Error("usage: picomap <info|load|log|test> [args...]")
		os.Exit(1)
	}
	if err != nil {
		slog.Error("fatal", "err", err)
		os.Exit(1)
	}
}

type deviceResult struct {
	dev  string
	info *client.ResponseInfo
	err  error
}

func printInfo(via string, from string, info *client.ResponseInfo) {
	slog.Info("device",
		"via", via,
		"from", from,
		"board_id", hex.EncodeToString(info.BoardID[:]),
		"mac", net.HardwareAddr(info.MAC[:]).String(),
		"ip", net.IP(info.IP[:]).String(),
		"firmware", info.FirmwareName)
}

func cmdInfo(args []string) error {
	fs := flag.NewFlagSet("info", flag.ExitOnError)
	udpAddr := fs.String("udp", "", "connect via UDP to this IP address")
	iface := fs.String("iface", "", "bind to this network interface (for broadcast)")
	fs.Parse(args)

	if *udpAddr == "" && *iface != "" {
		bcast, err := client.InterfaceBroadcast(*iface)
		if err != nil {
			return err
		}
		*udpAddr = bcast
	}

	if *udpAddr != "" {
		c, err := client.NewUDP(*udpAddr, *iface, 500*time.Millisecond)
		if err != nil {
			return err
		}
		defer c.Close()
		infos, err := c.InfoAll()
		if err != nil {
			return err
		}
		if len(infos) == 0 {
			return fmt.Errorf("no devices responded")
		}
		for _, r := range infos {
			printInfo(*udpAddr, r.From, r.Value)
		}
		return nil
	}

	devs, err := client.ListSerial()
	if err != nil {
		return err
	}
	if len(devs) == 0 {
		return fmt.Errorf("no devices found")
	}

	results := make([]deviceResult, len(devs))
	var wg sync.WaitGroup
	for i, dev := range devs {
		results[i].dev = dev
		wg.Go(func() {
			c, err := client.NewSerial(dev, 500*time.Millisecond)
			if err != nil {
				results[i].err = err
				return
			}
			info, err := c.Info()
			c.Close()
			if err != nil {
				results[i].err = err
				return
			}
			results[i].info = info
		})
	}
	wg.Wait()

	for _, r := range results {
		if r.err != nil {
			slog.Error("device error", "dev", r.dev, "err", r.err)
			continue
		}
		printInfo(r.dev, r.dev, r.info)
	}

	return nil
}

func boardSerial(id [8]byte) string {
	return fmt.Sprintf("%02X%02X%02X%02X%02X%02X%02X%02X",
		id[0], id[1], id[2], id[3], id[4], id[5], id[6], id[7])
}

type deviceInfo struct {
	dev    string
	serial string
	uf2    string
	name   string
}

func buildFirmware(buildDir string) error {
	slog.Info("configuring")
	cmake := exec.Command("cmake", "-S", filepath.Join(filepath.Dir(buildDir)), "-B", buildDir)
	cmake.Stdout = os.Stdout
	cmake.Stderr = os.Stderr
	if err := cmake.Run(); err != nil {
		return fmt.Errorf("cmake failed: %w", err)
	}

	slog.Info("building")
	cmd := exec.Command("make", "-C", buildDir)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("build failed: %w", err)
	}
	return nil
}

func cmdLog(_ []string) error {
	devs, err := client.ListSerial()
	if err != nil {
		return err
	}
	if len(devs) == 0 {
		return fmt.Errorf("no devices found")
	}
	for _, dev := range devs {
		log := slog.With("dev", dev)
		c, err := client.NewSerial(dev, 500*time.Millisecond)
		if err != nil {
			log.Error("connect error", "err", err)
			continue
		}
		resp, err := c.Log()
		c.Close()
		if err != nil {
			log.Error("log error", "err", err)
			continue
		}
		if len(resp.Entries) == 0 {
			log.Info("no debug messages")
			continue
		}
		for _, e := range resp.Entries {
			log.Info("dlog", "t_us", e.TimestampUS, "msg", e.Message)
		}
	}
	return nil
}

func cmdLoad(args []string) error {
	target := "all"
	if len(args) > 0 {
		target = args[0]
	}

	wd, err := os.Getwd()
	if err != nil {
		return err
	}
	buildDir := filepath.Join(wd, "firmware", "build")

	if err := buildFirmware(buildDir); err != nil {
		return err
	}

	devs, err := client.ListSerial()
	if err != nil {
		return err
	}

	allTargets := []struct {
		name string
		uf2  string
	}{
		{"picomap", filepath.Join(buildDir, "picomap.uf2")},
		{"picomap_test", filepath.Join(buildDir, "picomap_test.uf2")},
	}

	var targets []struct {
		name string
		uf2  string
	}
	switch target {
	case "all":
		targets = allTargets
	case "picomap":
		targets = allTargets[:1]
	case "picomap_test":
		targets = allTargets[1:]
	default:
		return fmt.Errorf("unknown target %q", target)
	}

	if len(devs) < len(targets) {
		return fmt.Errorf("need %d device(s), found %d", len(targets), len(devs))
	}

	devices := make([]deviceInfo, len(targets))
	errs := make([]error, len(targets))

	var wg sync.WaitGroup
	for i := range targets {
		log := slog.With("dev", devs[i])
		wg.Go(func() {
			c, err := client.NewSerial(devs[i], 500*time.Millisecond)
			if err != nil {
				errs[i] = err
				return
			}
			info, err := c.Info()
			c.Close()
			if err != nil {
				errs[i] = err
				return
			}
			devices[i] = deviceInfo{
				dev:    devs[i],
				serial: boardSerial(info.BoardID),
				uf2:    targets[i].uf2,
				name:   targets[i].name,
			}
			log.Info("got info", "serial", devices[i].serial, "firmware", info.FirmwareName)
		})
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			return fmt.Errorf("[%s] info: %w", devs[i], err)
		}
	}

	for i := range devices {
		log := slog.With("serial", devices[i].serial)
		wg.Go(func() {
			log.Info("sending PICOBOOT")
			c, err := client.NewSerial(devices[i].dev, 500*time.Millisecond)
			if err != nil {
				errs[i] = err
				return
			}
			err = c.PICOBOOT()
			c.Close()
			if err != nil {
				errs[i] = fmt.Errorf("PICOBOOT %s: %w", devices[i].serial, err)
				return
			}
			log.Info("PICOBOOT sent")
		})
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			return fmt.Errorf("[%s] %w", devices[i].serial, err)
		}
	}

	uf2s := make([]string, len(targets))
	for i := range targets {
		uf2s[i] = targets[i].uf2
	}

	for i := range devices {
		log := slog.With("serial", devices[i].serial)
		wg.Go(func() {
			log.Info("loading", "uf2", devices[i].name)
			errs[i] = picotool.Load(devices[i].uf2, devices[i].serial, 10*time.Second)
			if errs[i] == nil {
				log.Info("loaded", "uf2", devices[i].name)
			}
		})
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			return fmt.Errorf("[%s] load: %w", devices[i].serial, err)
		}
	}

	slog.Info("done")
	return nil
}

func findTestDevice() (string, error) {
	devs, err := client.ListSerial()
	if err != nil {
		return "", err
	}
	for _, dev := range devs {
		log := slog.With("dev", dev)
		c, err := client.NewSerial(dev, 500*time.Millisecond)
		if err != nil {
			log.Warn("connect error", "err", err)
			continue
		}
		info, err := c.Info()
		c.Close()
		if err != nil {
			log.Warn("info error", "err", err)
			continue
		}
		if info.FirmwareName == "picomap_test" {
			return dev, nil
		}
	}
	return "", fmt.Errorf("no picomap_test device found")
}

func cmdTestGroup(args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("usage: picomap test <list|run> [args...]")
	}
	switch args[0] {
	case "list":
		return cmdTestList(args[1:])
	case "run":
		return cmdTestRun(args[1:])
	default:
		return fmt.Errorf("usage: picomap test <list|run> [args...]")
	}
}

func cmdTestList(_ []string) error {
	dev, err := findTestDevice()
	if err != nil {
		return err
	}
	c, err := client.NewSerial(dev, 10*time.Second)
	if err != nil {
		return err
	}
	defer c.Close()

	result, err := c.ListTests()
	if err != nil {
		return fmt.Errorf("remote: %w", err)
	}
	for _, name := range result.Names {
		slog.Info("test", "dev", dev, "name", name)
	}
	return nil
}

func cmdTestRun(args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("usage: picomap test run <name>")
	}
	name := args[0]

	dev, err := findTestDevice()
	if err != nil {
		return err
	}

	log := slog.With("dev", dev)
	log.Info("running test", "name", name)

	c, err := client.NewSerial(dev, 10*time.Second)
	if err != nil {
		return err
	}
	defer c.Close()

	result, err := c.Test(name)
	if err != nil {
		return fmt.Errorf("remote: %w", err)
	}

	for _, msg := range result.Messages {
		log.Info("remote", "msg", msg)
	}

	if result.Pass {
		log.Info("PASS")
	} else {
		log.Error("FAIL")
		os.Exit(1)
	}
	return nil
}