picomap/cmd/picomap/main.go

package main

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

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

type target struct {
	name   string
	client *client.Client
}

type targetFlags struct {
	usb   string
	udp   string
	iface string
}

func addTargetFlags(fs *flag.FlagSet) *targetFlags {
	tf := &targetFlags{}
	fs.StringVar(&tf.usb, "usb", "", "comma-separated USB serial devices")
	fs.StringVar(&tf.udp, "udp", "", "comma-separated UDP IP addresses")
	fs.StringVar(&tf.iface, "iface", "", "network interface for multicast discovery")
	return tf
}

func (tf *targetFlags) connect(timeout time.Duration) ([]target, error) {
	var targets []target

	if tf.usb != "" {
		for _, dev := range strings.Split(tf.usb, ",") {
			dev = strings.TrimSpace(dev)
			c, err := client.NewSerial(dev, timeout)
			if err != nil {
				return nil, fmt.Errorf("usb %s: %w", dev, err)
			}
			targets = append(targets, target{dev, c})
		}
	}

	if tf.udp != "" {
		for _, addr := range strings.Split(tf.udp, ",") {
			addr = strings.TrimSpace(addr)
			c, err := client.NewUDP(addr, tf.iface, timeout)
			if err != nil {
				return nil, fmt.Errorf("udp %s: %w", addr, err)
			}
			targets = append(targets, target{addr, c})
		}
	}

	if tf.iface != "" && tf.udp == "" {
		bcast, err := client.InterfaceBroadcast(tf.iface)
		if err != nil {
			return nil, err
		}
		c, err := client.NewUDP(bcast, tf.iface, timeout)
		if err != nil {
			return nil, err
		}
		infos, err := c.InfoAll()
		c.Close()
		if err != nil {
			return nil, fmt.Errorf("discovery: %w", err)
		}
		for _, r := range infos {
			uc, err := client.NewUDP(r.From, tf.iface, timeout)
			if err != nil {
				return nil, fmt.Errorf("udp %s: %w", r.From, err)
			}
			targets = append(targets, target{r.From, uc})
		}
	}

	if tf.usb == "" && tf.udp == "" && tf.iface == "" {
		devs, err := client.ListSerial()
		if err != nil {
			return nil, err
		}
		for _, dev := range devs {
			c, err := client.NewSerial(dev, timeout)
			if err != nil {
				slog.Warn("connect error", "dev", dev, "err", err)
				continue
			}
			targets = append(targets, target{dev, c})
		}
	}

	if len(targets) == 0 {
		return nil, fmt.Errorf("no devices found")
	}
	return targets, nil
}

func closeTargets(targets []target) {
	for _, t := range targets {
		t.client.Close()
	}
}

func main() {
	if len(os.Args) < 2 {
		slog.Error("usage: picomap <info|load|log|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)
	}
}

func printInfo(via string, info *client.ResponseInfo) {
	slog.Info("device",
		"via", via,
		"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)
	tf := addTargetFlags(fs)
	fs.Parse(args)

	targets, err := tf.connect(500*time.Millisecond)
	if err != nil {
		return err
	}
	defer closeTargets(targets)

	for _, t := range targets {
		info, err := t.client.Info()
		if err != nil {
			slog.Error("info error", "via", t.name, "err", err)
			continue
		}
		printInfo(t.name, info)
	}
	return nil
}

func cmdLog(args []string) error {
	fs := flag.NewFlagSet("log", flag.ExitOnError)
	tf := addTargetFlags(fs)
	fs.Parse(args)

	targets, err := tf.connect(500*time.Millisecond)
	if err != nil {
		return err
	}
	defer closeTargets(targets)

	for _, t := range targets {
		log := slog.With("dev", t.name)
		resp, err := t.client.Log()
		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 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])
}

func runCmd(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	out, err := cmd.CombinedOutput()
	if len(out) > 0 {
		for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
			slog.Info(name, "msg", line)
		}
	}
	if err != nil {
		return fmt.Errorf("%s failed: %w", name, err)
	}
	return nil
}

func buildFirmware(buildDir string) error {
	srcDir := filepath.Dir(buildDir)
	if err := runCmd("cmake", "-S", srcDir, "-B", buildDir); err != nil {
		return err
	}
	return runCmd("cmake", "--build", buildDir)
}

func cmdLoad(args []string) error {
	fs := flag.NewFlagSet("load", flag.ExitOnError)
	tf := addTargetFlags(fs)
	dryRun := fs.Bool("dry-run", false, "parse UF2 and log operations without flashing")
	fs.Parse(args)
	loadTarget := "all"
	if fs.NArg() > 0 {
		loadTarget = fs.Arg(0)
	}

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

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

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

	var fwTargets []firmwareTarget
	switch loadTarget {
	case "all":
		fwTargets = allTargets
	case "picomap":
		fwTargets = allTargets[:1]
	case "picomap_test":
		fwTargets = allTargets[1:]
	default:
		return fmt.Errorf("unknown target %q", loadTarget)
	}

	targets, err := tf.connect(5*time.Second)
	if err != nil {
		return err
	}
	defer closeTargets(targets)

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

	type deviceInfo struct {
		target target
		serial string
		fw     firmwareTarget
	}

	devices := make([]deviceInfo, len(fwTargets))
	for i, t := range targets[:len(fwTargets)] {
		info, err := t.client.Info()
		if err != nil {
			return fmt.Errorf("[%s] info: %w", t.name, err)
		}
		devices[i] = deviceInfo{
			target: t,
			serial: boardSerial(info.BoardID),
			fw:     fwTargets[i],
		}
		slog.Info("got info", "dev", t.name, "serial", devices[i].serial, "firmware", info.FirmwareName)
	}

	errs := make([]error, len(devices))
	var wg sync.WaitGroup
	for i := range devices {
		log := slog.With("dev", devices[i].target.name, "serial", devices[i].serial)
		wg.Go(func() {
			log.Info("flashing", "uf2", devices[i].fw.name)
			errs[i] = flashDevice(devices[i].target.client, devices[i].fw.uf2, *dryRun, log)
			if errs[i] == nil {
				log.Info("flashed", "uf2", devices[i].fw.name)
			}
		})
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			return fmt.Errorf("[%s] flash: %w", devices[i].serial, err)
		}
	}

	slog.Info("done")
	return nil
}

func flashDevice(c *client.Client, uf2Path string, dryRun bool, log *slog.Logger) error {
	blocks, err := uf2.Parse(uf2Path)
	if err != nil {
		return fmt.Errorf("parse uf2: %w", err)
	}

	log.Info("parsed uf2", "blocks", len(blocks))

	const sectorSize = 4096

	if dryRun {
		erased := make(map[uint32]bool)
		for _, b := range blocks {
			sector := b.Addr &^ (sectorSize - 1)
			if !erased[sector] {
				log.Info("erasing", "addr", fmt.Sprintf("%08x", sector))
				erased[sector] = true
			}
			log.Info("writing", "addr", fmt.Sprintf("%08x", b.Addr), "len", len(b.Data))
		}
		return nil
	}

	erased := make(map[uint32]bool)
	for _, b := range blocks {
		sector := b.Addr &^ (sectorSize - 1)
		if !erased[sector] {
			if err := c.FlashErase(sector, sectorSize); err != nil {
				return fmt.Errorf("erase %08x: %w", sector, err)
			}
			erased[sector] = true
		}
		if err := c.FlashWrite(b.Addr, b.Data); err != nil {
			return fmt.Errorf("write %08x: %w", b.Addr, err)
		}
	}

	log.Info("rebooting")
	c.Reboot()
	return nil
}

func findTestDevice(args []string) (target, error) {
	fs := flag.NewFlagSet("test", flag.ExitOnError)
	tf := addTargetFlags(fs)
	fs.Parse(args)

	targets, err := tf.connect(10*time.Second)
	if err != nil {
		return target{}, err
	}

	for _, t := range targets {
		info, err := t.client.Info()
		if err != nil {
			t.client.Close()
			continue
		}
		if info.FirmwareName == "picomap_test" {
			for _, other := range targets {
				if other.name != t.name {
					other.client.Close()
				}
			}
			return t, nil
		}
		t.client.Close()
	}
	return target{}, fmt.Errorf("no picomap_test device found")
}

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

func cmdTestList(args []string) error {
	t, err := findTestDevice(args)
	if err != nil {
		return err
	}
	defer t.client.Close()

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

func cmdTestAll(args []string) error {
	t, err := findTestDevice(args)
	if err != nil {
		return err
	}
	defer t.client.Close()

	list, err := t.client.ListTests()
	if err != nil {
		return fmt.Errorf("remote: %w", err)
	}

	log := slog.With("dev", t.name)
	failed := 0
	for _, name := range list.Names {
		log.Info("running test", "name", name)
		result, err := t.client.Test(name)
		if err != nil {
			log.Error("error", "name", name, "err", err)
			failed++
			continue
		}
		for _, msg := range result.Messages {
			log.Info("remote", "name", name, "msg", msg)
		}
		if result.Pass {
			log.Info("PASS", "name", name)
		} else {
			log.Error("FAIL", "name", name)
			failed++
		}
	}
	if failed > 0 {
		log.Error("tests failed", "count", failed)
		os.Exit(1)
	}
	log.Info("all tests passed", "count", len(list.Names))
	return nil
}

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

	t, err := findTestDevice(args[1:])
	if err != nil {
		return err
	}
	defer t.client.Close()

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

	result, err := t.client.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
}