spoolweight/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"log/slog"
	"math"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/makiuchi-d/gozxing"
	"github.com/makiuchi-d/gozxing/qrcode"
	"spoolweight/claude"
	"spoolweight/spooldb"
)

// weightBreakdown is the spool/filament/total split, in grams.
type weightBreakdown struct {
	Spool    *float64 `json:"spool,omitempty"`
	Filament *float64 `json:"filament,omitempty"`
	Total    *float64 `json:"total,omitempty"`
}

type result struct {
	Image            string           `json:"image"`
	SpoolID          string           `json:"spool_id,omitempty"`
	URL              string           `json:"url,omitempty"`
	Location         string           `json:"location,omitempty"`
	PreviousWeight   *weightBreakdown `json:"previous_weight,omitempty"` // as read from spooldb
	NewWeight        *weightBreakdown `json:"new_weight,omitempty"`      // from the measured photo
	Unit             string           `json:"unit,omitempty"`
	Confidence       *float64         `json:"confidence,omitempty"`
	VoteWeights      []float64        `json:"vote_weights,omitempty"`      // each LLM vote's total reading
	ModelConfidences []string         `json:"model_confidences,omitempty"` // each LLM vote's self-rating
	Updated          *bool            `json:"updated,omitempty"`           // set on success: did we write to spooldb
	Error            string           `json:"error,omitempty"`
}

func ptr(v float64) *float64 { return &v }

// spoolURLPrefix is the only QR payload we accept. Anything else (a different
// domain, or a /f/ filament link rather than a /s/ spool link) is rejected.
const spoolURLPrefix = "https://spooldb.com/s/"

func main() {
	verbose := flag.Bool("v", false, "verbose: log each event's start/end with timings to stderr")
	dryRun := flag.Bool("n", false, "dry run: read and report, but do not write weight changes back to the site")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: %s [-v] [-n] <image>...\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(2)
	}

	level := slog.LevelInfo
	if *verbose {
		level = slog.LevelDebug
	}
	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level})))

	auth, err := claude.ResolveAuth()
	if err != nil {
		fail("resolve credentials: %v", err)
	}
	slog.Debug("claude auth", "mechanism", auth.Name())

	sp := &spoolSync{}
	defer sp.close()

	results := make([]result, 0, flag.NArg())
	for _, path := range flag.Args() {
		results = append(results, processImage(path, auth, sp, *dryRun))
	}

	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	enc.Encode(results)
}

// spoolSync lazily logs into spooldb on first use, so its (slow) browser login
// overlaps with the first image's weight read rather than blocking it. Every
// site interaction is best-effort: failures are logged and the corresponding
// output fields are simply omitted.
type spoolSync struct {
	once   sync.Once
	client *spooldb.Client
	err    error
}

func (s *spoolSync) login() (*spooldb.Client, error) {
	s.once.Do(func() {
		user, pass := os.Getenv("SPOOLDB_USER"), os.Getenv("SPOOLDB_PASS")
		if user == "" || pass == "" {
			s.err = fmt.Errorf("SPOOLDB_USER/SPOOLDB_PASS not set")
			return
		}
		c, err := spooldb.New()
		if err != nil {
			s.err = err
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
		defer cancel()
		if err := c.Login(ctx, user, pass); err != nil {
			c.Close()
			s.err = fmt.Errorf("login: %w", err)
			return
		}
		s.client = c
	})
	return s.client, s.err
}

// info reads the spool's current location and weights, or nil if unavailable.
func (s *spoolSync) info(spoolID string) *spooldb.SpoolInfo {
	client, err := s.login()
	if err != nil {
		slog.Warn("spooldb: skipping spool lookup", "err", err)
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	info, err := client.SpoolInfo(ctx, spoolID)
	if err != nil {
		slog.Warn("spooldb: lookup failed", "spool", spoolID, "err", err)
		return nil
	}
	return &info
}

// setTotal writes a new total (with-spool) weight back to the site.
func (s *spoolSync) setTotal(spoolID string, grams float64) bool {
	client, err := s.login()
	if err != nil {
		slog.Warn("spooldb: skipping weight update", "err", err)
		return false
	}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	if err := client.SetTotalWeight(ctx, spoolID, grams); err != nil {
		slog.Warn("spooldb: weight update failed", "spool", spoolID, "err", err)
		return false
	}
	return true
}

func (s *spoolSync) close() {
	if s.client != nil {
		s.client.Close()
	}
}

// processImage reads one photo, capturing any failure in the result's Error
// field so a single bad image doesn't abort the whole batch.
func processImage(path string, auth claude.Auth, sp *spoolSync, dryRun bool) result {
	defer trace("image " + path)()
	r := result{Image: path}

	img, err := loadImage(path)
	if err != nil {
		r.Error = fmt.Sprintf("load image: %v", err)
		return r
	}

	url, err := decodeQR(img)
	if err != nil {
		r.Error = fmt.Sprintf("read QR code: %v", err)
		return r
	}
	if !strings.HasPrefix(url, spoolURLPrefix) {
		r.Error = fmt.Sprintf("not a spool QR code (expected %s...): %s", spoolURLPrefix, url)
		return r
	}
	r.SpoolID = spoolID(url)
	r.URL = url

	// Read the spool's current location/weight from the site concurrently with
	// the (slower) weight read off the photo.
	var info *spooldb.SpoolInfo
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		info = sp.info(r.SpoolID)
	}()

	reading, err := claude.ReadWeight(img, auth)
	wg.Wait()
	if info != nil {
		r.Location = info.Location
		r.PreviousWeight = &weightBreakdown{
			Spool:    ptr(info.EmptySpoolGrams),
			Filament: ptr(info.RemainingGrams),
			Total:    ptr(info.TotalGrams),
		}
	}

	if err != nil {
		r.Error = fmt.Sprintf("read weight: %v", err)
		return r
	}
	if reading.Unit == "lb" || reading.Unit == "oz" {
		r.Error = fmt.Sprintf("scale is set to imperial units (%s); switch it to grams", reading.Unit)
		return r
	}
	r.Unit = reading.Unit
	r.Confidence = ptr(reading.Confidence)
	r.VoteWeights = reading.Weights
	r.ModelConfidences = reading.ModelConfidences

	// The photo gives total weight (filament + spool). When the empty-spool
	// weight is known (from spooldb) we can also report the filament split.
	r.NewWeight = &weightBreakdown{Total: ptr(reading.Weight)}
	if info != nil {
		r.NewWeight.Spool = ptr(info.EmptySpoolGrams)
		r.NewWeight.Filament = ptr(reading.Weight - info.EmptySpoolGrams)
	}

	// Write the measured total back when it differs from the site's total — the
	// site recomputes remaining filament using the empty-spool weight.
	updated := false
	if info != nil && math.Abs(reading.Weight-info.TotalGrams) >= 1 {
		newRemaining := reading.Weight - info.EmptySpoolGrams
		if dryRun {
			slog.Info("spooldb: would update", "spool", r.SpoolID,
				"remaining_from", info.RemainingGrams, "remaining_to", newRemaining, "measured_total", reading.Weight)
		} else if sp.setTotal(r.SpoolID, reading.Weight) {
			updated = true
			slog.Info("spooldb: updated", "spool", r.SpoolID,
				"remaining_from", info.RemainingGrams, "remaining_to", newRemaining, "measured_total", reading.Weight)
		}
	}
	r.Updated = &updated
	return r
}

func fail(format string, a ...interface{}) {
	slog.Error(fmt.Sprintf(format, a...))
	os.Exit(1)
}

// trace logs the start of an operation and, via the returned func, its end and
// duration. It is a no-op above debug level, so it only appears with -v.
func trace(op string) func() {
	t0 := time.Now()
	slog.Debug("begin", "op", op)
	return func() { slog.Debug("end", "op", op, "dur", time.Since(t0)) }
}

func loadImage(path string) (image.Image, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	img, _, err := image.Decode(f)
	return img, err
}

// spoolID returns the last non-empty path segment of the QR URL.
func spoolID(url string) string {
	s := strings.TrimRight(url, "/")
	if i := strings.LastIndex(s, "/"); i >= 0 {
		return s[i+1:]
	}
	return s
}

func decodeQR(img image.Image) (string, error) {
	bmp, err := gozxing.NewBinaryBitmapFromImage(img)
	if err != nil {
		return "", err
	}
	res, err := qrcode.NewQRCodeReader().Decode(bmp, map[gozxing.DecodeHintType]interface{}{
		gozxing.DecodeHintType_TRY_HARDER: true,
	})
	if err != nil {
		return "", err
	}
	return res.GetText(), nil
}