ImaGen/cmd/imagen/usage.go

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
	"text/tabwriter"
	"time"

	"mgit.msbls.de/m/ImaGen/internal/usage"
)

// runUsage handles `imagen usage [--since DATE]`. Reads mai.imagen_usage
// via Supabase REST and prints a tab-aligned table grouped by week +
// backend + model + caller, with totals at the bottom.
func runUsage(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("usage", flag.ContinueOnError)
	var (
		since string
		raw   bool
	)
	fs.StringVar(&since, "since", "", "ISO date (YYYY-MM-DD) — only rows on/after this UTC date")
	fs.BoolVar(&raw, "raw", false, "print one line per row instead of grouped")
	fs.Usage = func() {
		fmt.Fprintln(fs.Output(), "Usage: imagen usage [--since YYYY-MM-DD] [--raw]")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
	}

	var sinceT time.Time
	if since != "" {
		t, err := time.Parse("2006-01-02", since)
		if err != nil {
			return userErr("--since must be YYYY-MM-DD: %v", err)
		}
		sinceT = t
	}

	sink, ok := usage.NewSupabaseSinkFromEnv()
	if !ok {
		return userErr("SUPABASE_URL and SUPABASE_SERVICE_KEY (or MAI_SUPABASE_KEY) must be set to read mai.imagen_usage")
	}
	rows, err := sink.Query(ctx, sinceT)
	if err != nil {
		return err
	}

	if raw {
		printRawRows(rows)
		return nil
	}
	printGroupedRows(rows)
	return nil
}

func printRawRows(rows []usage.Row) {
	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintln(tw, "TIME\tBACKEND\tMODEL\tCALLER\tLATENCY_MS\tCOST_USD")
	var totalCost float64
	for _, r := range rows {
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%s\n",
			r.CreatedAt.Local().Format("2006-01-02 15:04"),
			r.Backend,
			r.Model,
			derefString(r.Caller),
			intOrDash(r.LatencyMs),
			costOrDash(r.CostUSDEstimate),
		)
		if r.CostUSDEstimate != nil {
			totalCost += *r.CostUSDEstimate
		}
	}
	fmt.Fprintf(tw, "\t\t\t\t%d rows\t%.4f USD\n", len(rows), totalCost)
	_ = tw.Flush()
}

type group struct {
	week    string
	backend string
	model   string
	caller  string
	count   int
	cost    float64
	costSet bool
}

type groupKey struct {
	week, backend, model, caller string
}

func printGroupedRows(rows []usage.Row) {
	groups := map[groupKey]*group{}
	for _, r := range rows {
		caller := derefString(r.Caller)
		k := groupKey{
			week:    weekStart(r.CreatedAt).Format("2006-01-02"),
			backend: r.Backend,
			model:   r.Model,
			caller:  caller,
		}
		g, ok := groups[k]
		if !ok {
			g = &group{week: k.week, backend: r.Backend, model: r.Model, caller: caller}
			groups[k] = g
		}
		g.count++
		if r.CostUSDEstimate != nil {
			g.cost += *r.CostUSDEstimate
			g.costSet = true
		}
	}

	keys := make([]groupKey, 0, len(groups))
	for k := range groups {
		keys = append(keys, k)
	}
	sort.Slice(keys, func(i, j int) bool {
		if keys[i].week != keys[j].week {
			return keys[i].week > keys[j].week // newest first
		}
		if keys[i].backend != keys[j].backend {
			return keys[i].backend < keys[j].backend
		}
		if keys[i].model != keys[j].model {
			return keys[i].model < keys[j].model
		}
		return keys[i].caller < keys[j].caller
	})

	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintln(tw, "WEEK_OF\tBACKEND\tMODEL\tCALLER\tCOUNT\tCOST_USD")
	var totalCount int
	var totalCost float64
	for _, k := range keys {
		g := groups[k]
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%d\t%s\n",
			g.week, g.backend, g.model, g.caller, g.count, costStr(g.cost, g.costSet),
		)
		totalCount += g.count
		totalCost += g.cost
	}
	fmt.Fprintf(tw, "\t\t\tTOTAL\t%d\t%.4f USD\n", totalCount, totalCost)
	_ = tw.Flush()
}

// weekStart returns the Monday of the week containing t (UTC).
func weekStart(t time.Time) time.Time {
	t = t.UTC()
	wd := int(t.Weekday())
	if wd == 0 {
		wd = 7 // shift Sunday to end-of-week
	}
	delta := time.Duration(wd-1) * -24 * time.Hour
	d := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
	return d.Add(delta)
}

func derefString(s *string) string {
	if s == nil {
		return ""
	}
	return *s
}

func intOrDash(p *int) string {
	if p == nil {
		return "-"
	}
	return fmt.Sprintf("%d", *p)
}

func costOrDash(p *float64) string {
	if p == nil {
		return "-"
	}
	return fmt.Sprintf("%.4f", *p)
}

func costStr(v float64, set bool) string {
	if !set {
		return "-"
	}
	return strings.TrimSpace(fmt.Sprintf("%.4f", v))
}