zenmux-usage-cli/cmd/zenmux-usage/main.go

// Command zenmux-usage prints ZenMux subscription quota windows and the
// USD value of tokens consumed for one or more configured accounts.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/fatih/color"
	"github.com/mattn/go-isatty"

	"github.com/kotoyuuko/zenmux-usage-cli/internal/api"
	"github.com/kotoyuuko/zenmux-usage-cli/internal/config"
	"github.com/kotoyuuko/zenmux-usage-cli/internal/render"
)

// version is overridable at build time: -ldflags "-X main.version=v0.1.0".
var version = "dev"

// Exit codes. See design §9 and specs/subscription-usage/spec.md.
const (
	exitOK          = 0
	exitGeneric     = 1
	exitUsage       = 2
	exitNoAccount   = 3
	exitAuth        = 4
	exitRateLimited = 5
	exitNetwork     = 6
	exitConfigParse = 7
)

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
}

type opts struct {
	account  string
	config   string
	apiKey   string
	json     bool
	noColor  bool
	timeout  time.Duration
	version  bool
	// Internal: set when --config was explicitly provided (vs defaulted).
	configExplicit bool
}

func run(args []string, stdout, stderr io.Writer) int {
	o, proceed, code := parseFlags(args, stderr)
	if !proceed {
		return code
	}
	if o.version {
		fmt.Fprintf(stdout, "zenmux-usage %s\n", version)
		return exitOK
	}

	// Resolve accounts to fetch.
	accounts, cfgLoaded, code := resolveAccounts(o, stderr)
	if code != exitOK {
		return code
	}

	// Permissions warning only fires when a config file was actually loaded.
	if cfgLoaded != "" {
		config.WarnIfLooseMode(cfgLoaded, stderr)
	}

	// Fetch all accounts sequentially.
	client := api.NewClient(o.timeout)
	if base := os.Getenv("ZENMUX_BASE_URL"); base != "" {
		client.BaseURL = base
	}
	results, rawBodies := fetchAll(client, accounts, o.timeout)

	// Render output.
	if o.json {
		if len(results) == 1 && results[0].Err == nil {
			if err := render.RenderJSONSingle(stdout, rawBodies[0]); err != nil {
				fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
				return exitGeneric
			}
		} else {
			if err := render.RenderJSONMulti(stdout, results); err != nil {
				fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
				return exitGeneric
			}
		}
	} else {
		useColor := shouldUseColor(o, stdout)
		render.RenderAll(stdout, results, useColor)
	}

	return computeExitCode(results)
}

// parseFlags parses args into opts. Returns (opts, proceed, code):
// proceed=true means run() should continue; proceed=false means return code
// immediately (used for --help, --version, flag errors).
func parseFlags(args []string, stderr io.Writer) (opts, bool, int) {
	var o opts
	fs := flag.NewFlagSet("zenmux-usage", flag.ContinueOnError)
	fs.SetOutput(stderr)
	fs.StringVar(&o.account, "account", "", "filter to a single account by name (from config)")
	fs.StringVar(&o.config, "config", "", "path to config file (default: $XDG_CONFIG_HOME/zenmux-usage/config.yaml)")
	fs.StringVar(&o.apiKey, "api-key", "", "use this API key directly; bypass config file")
	fs.BoolVar(&o.json, "json", false, "emit JSON on stdout instead of the human layout")
	fs.BoolVar(&o.noColor, "no-color", false, "disable ANSI colors in human output")
	fs.DurationVar(&o.timeout, "timeout", 10*time.Second, "HTTP timeout per account")
	fs.BoolVar(&o.version, "version", false, "print version and exit")
	fs.Usage = func() {
		fmt.Fprintf(stderr, "Usage: %s [flags]\n\nFlags:\n", fs.Name())
		fs.PrintDefaults()
	}

	if err := fs.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return o, false, exitOK
		}
		return o, false, exitUsage
	}
	if fs.NArg() > 0 {
		fmt.Fprintf(stderr, "zenmux-usage: unexpected positional arguments: %v\n", fs.Args())
		return o, false, exitUsage
	}
	fs.Visit(func(f *flag.Flag) {
		if f.Name == "config" {
			o.configExplicit = true
		}
	})
	if o.timeout <= 0 {
		fmt.Fprintln(stderr, "zenmux-usage: --timeout must be positive")
		return o, false, exitUsage
	}
	return o, true, exitOK
}

// resolveAccounts loads the config (if any) and produces the list of accounts
// to fetch. Returns the list, the loaded-config path (for permission warning;
// empty if no config file was read), and an exit code (exitOK on success).
func resolveAccounts(o opts, stderr io.Writer) ([]config.Account, string, int) {
	// --api-key short-circuits config loading entirely.
	if o.apiKey != "" {
		accs, err := config.Resolve(nil, config.ResolveFlags{APIKey: o.apiKey})
		if err != nil {
			fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
			return nil, "", exitGeneric
		}
		return accs, "", exitOK
	}

	cfgPath := o.config
	if cfgPath == "" {
		cfgPath = config.DefaultPath()
	}
	var cfg *config.Config
	loadedPath := ""
	if cfgPath != "" {
		c, err := config.Load(cfgPath, stderr)
		switch {
		case err == nil:
			cfg = c
			loadedPath = cfgPath
		case errors.Is(err, config.ErrParse):
			fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
			return nil, "", exitConfigParse
		case os.IsNotExist(err):
			// Missing file: only fatal if caller explicitly asked for that path.
			if o.configExplicit {
				fmt.Fprintf(stderr, "zenmux-usage: config file not found: %s\n", cfgPath)
				return nil, "", exitNoAccount
			}
			// Default path missing → fall through to env-var fallback.
		default:
			fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
			return nil, "", exitGeneric
		}
	}

	envKey := os.Getenv("ZENMUX_MANAGEMENT_API_KEY")
	accs, err := config.Resolve(cfg, config.ResolveFlags{
		AccountName: o.account,
		EnvAPIKey:   envKey,
	})
	if err != nil {
		switch {
		case errors.Is(err, config.ErrAccountNotFound):
			fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
			return nil, "", exitUsage
		case errors.Is(err, config.ErrNoAccount):
			fmt.Fprintf(stderr,
				"zenmux-usage: no account available. Create %s or set ZENMUX_MANAGEMENT_API_KEY.\n",
				config.DefaultPath())
			return nil, "", exitNoAccount
		default:
			fmt.Fprintf(stderr, "zenmux-usage: %v\n", err)
			return nil, "", exitGeneric
		}
	}
	return accs, loadedPath, exitOK
}

// fetchAll runs each account's request sequentially. Errors are attached
// per-result rather than aborting the run.
func fetchAll(client *api.Client, accounts []config.Account, timeout time.Duration) ([]render.AccountResult, [][]byte) {
	results := make([]render.AccountResult, 0, len(accounts))
	raws := make([][]byte, 0, len(accounts))
	for _, acc := range accounts {
		ctx, cancel := context.WithTimeout(context.Background(), timeout+time.Second)
		resp, raw, err := client.FetchSubscriptionDetail(ctx, acc.APIKey)
		cancel()
		results = append(results, render.AccountResult{
			Name:     acc.Name,
			Response: resp,
			Err:      err,
		})
		raws = append(raws, raw)
	}
	return results, raws
}

// computeExitCode applies the rules from design §9:
//   - all succeeded → 0
//   - all failed with the same cause → specific code (auth/rate/network)
//   - mixed success+failure, or mixed causes → 1
func computeExitCode(results []render.AccountResult) int {
	if len(results) == 0 {
		return exitGeneric
	}
	successes := 0
	causeCounts := map[int]int{}
	for _, r := range results {
		if r.Err == nil {
			successes++
			continue
		}
		causeCounts[classify(r.Err)]++
	}
	if successes == len(results) {
		return exitOK
	}
	if successes > 0 {
		return exitGeneric // mixed success + failure
	}
	// All failed. Check single-cause.
	if len(causeCounts) == 1 {
		for code := range causeCounts {
			return code
		}
	}
	return exitGeneric
}

func classify(err error) int {
	switch {
	case errors.Is(err, api.ErrUnauthorized):
		return exitAuth
	case errors.Is(err, api.ErrRateLimited):
		return exitRateLimited
	case errors.Is(err, api.ErrTimeout):
		return exitNetwork
	default:
		return exitGeneric
	}
}

// shouldUseColor resolves the color decision: user flag wins, then NO_COLOR,
// then stdout TTY detection.
func shouldUseColor(o opts, stdout io.Writer) bool {
	if o.noColor {
		return false
	}
	if os.Getenv("NO_COLOR") != "" {
		return false
	}
	// Prefer the file-descriptor test for the actual stdout.
	if f, ok := stdout.(*os.File); ok {
		return isatty.IsTerminal(f.Fd()) || isatty.IsCygwinTerminal(f.Fd())
	}
	// Non-*os.File writer (e.g. test buffer) → stay monochrome.
	_ = color.NoColor
	return false
}