ImaGen/cmd/imagen/generate.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"mgit.msbls.de/m/ImaGen/internal/backend"
	"mgit.msbls.de/m/ImaGen/internal/cloud"
	"mgit.msbls.de/m/ImaGen/internal/config"
	"mgit.msbls.de/m/ImaGen/internal/output"
	"mgit.msbls.de/m/ImaGen/internal/preview"
	"mgit.msbls.de/m/ImaGen/internal/prompt"
	"mgit.msbls.de/m/ImaGen/internal/usage"
)

func runGenerate(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("generate", flag.ContinueOnError)
	var (
		backendName  string
		size         string
		outPath      string
		seed         int64
		steps        int
		style        string
		negative     string
		configPath   string
		noSidecar    bool
		previewOn    bool
		previewOff   bool
		noCloud      bool
	)
	fs.StringVar(&backendName, "backend", "", "backend instance name (default: config.default_backend)")
	fs.StringVar(&size, "size", "1024x1024", "WxH, e.g. 1024x1024")
	fs.StringVar(&outPath, "output", "", "explicit output path (overrides config naming template)")
	fs.Int64Var(&seed, "seed", 0, "deterministic seed (0 = backend default)")
	fs.IntVar(&steps, "steps", 0, "diffusion steps (0 = backend default)")
	fs.StringVar(&style, "style", "", "style preset name (see imagen config init for the list)")
	fs.StringVar(&negative, "negative", "", "negative prompt (ignored by backends that don't support it)")
	fs.StringVar(&configPath, "config", "", "config file path (default: ~/.config/imagen.yaml)")
	fs.BoolVar(&noSidecar, "no-sidecar", false, "skip the JSON sidecar even if config enables it")
	fs.BoolVar(&previewOn, "preview", false, "force tmux preview window on (errors outside $TMUX)")
	fs.BoolVar(&previewOff, "no-preview", false, "skip the tmux preview window")
	fs.BoolVar(&noCloud, "no-cloud", false, "skip Supabase upload + imagen.images insert for this generation")
	fs.Usage = func() {
		fmt.Fprintln(fs.Output(), `Usage: imagen generate "<prompt>" [flags]`)
		fs.PrintDefaults()
	}
	// stdlib flag stops parsing at the first non-flag arg, so split the
	// prompt (leading positional args) from the flags ourselves before parsing.
	leadingPositional, flagArgs := splitLeadingPositional(args)
	if err := fs.Parse(flagArgs); err != nil {
		return err
	}
	positional := append(leadingPositional, fs.Args()...)
	if len(positional) == 0 {
		fs.Usage()
		return userErr("missing prompt")
	}
	rawPrompt := strings.Join(positional, " ")

	w, h, err := parseSize(size)
	if err != nil {
		return userErr("bad --size: %v", err)
	}

	cfg, cfgErr := config.Load(configPath)
	if cfgErr != nil && !os.IsNotExist(cfgErr) {
		return cfgErr
	}

	if backendName == "" {
		if cfg != nil {
			backendName = cfg.DefaultBackend
		}
	}
	if backendName == "" {
		return userErr("no --backend given and no default_backend in config")
	}

	be, err := buildBackend(cfg, backendName)
	if err != nil {
		return err
	}
	attachUsageSink(be)

	finalPrompt, err := prompt.Apply(rawPrompt, style)
	if err != nil {
		return userErr("%v", err)
	}

	req := backend.Request{
		Prompt:         finalPrompt,
		NegativePrompt: negative,
		Width:          w,
		Height:         h,
		Steps:          steps,
		Seed:           seed,
		Style:          style,
	}
	res, err := be.Generate(ctx, req)
	if err != nil {
		return fmt.Errorf("backend %q: %w", backendName, err)
	}
	defer res.ImageReader.Close()

	writer := buildWriter(cfg, noSidecar)
	in := output.Inputs{
		Prompt:   rawPrompt,
		Backend:  be.Name(),
		Seed:     seedFromMetadata(res.Metadata, seed),
		Ext:      extFromMime(res.MimeType),
		Metadata: res.Metadata,
	}
	var paths *output.Outputs
	if outPath != "" {
		paths, err = writer.WriteToPath(res.ImageReader, outPath, in)
	} else {
		paths, err = writer.Write(res.ImageReader, in)
	}
	if err != nil {
		return err
	}
	fmt.Println(paths.ImagePath)
	if paths.SidecarPath != "" {
		fmt.Fprintln(os.Stderr, "sidecar:", paths.SidecarPath)
	}

	if result, err := maybeCloudSync(ctx, cfg, noCloud, "", paths, in, res, w, h); err != nil {
		// cloud-sync failures are warnings — the image already wrote.
		fmt.Fprintln(os.Stderr, "imagen: cloud sync:", err)
	} else if result != nil && result.ImageID != "" {
		fmt.Fprintf(os.Stderr, "cloud: imagen.images.id=%s storage_path=%s\n", result.ImageID, result.StoragePath)
	}

	if err := maybePreview(cfg, previewOn, previewOff, paths.ImagePath, rawPrompt); err != nil {
		// preview failures are warnings — the image already wrote.
		fmt.Fprintln(os.Stderr, "imagen: preview:", err)
	}
	return nil
}

// resolveCloudSyncMode applies the precedence chain config -> env -> flag.
// Flags win, env beats config, config beats the implicit auto default.
// Mirrors resolvePreviewMode shape.
func resolveCloudSyncMode(cfg *config.Config, noCloudFlag bool, env string) (string, error) {
	mode := "auto"
	if cfg != nil && cfg.Output.CloudSync != "" {
		mode = cfg.Output.CloudSync
	}
	if env != "" {
		switch env {
		case "auto", "on", "off":
			mode = env
		default:
			return "", fmt.Errorf("$IMAGEN_CLOUD_SYNC = %q (must be auto|on|off)", env)
		}
	}
	if noCloudFlag {
		mode = "off"
	}
	return mode, nil
}

// maybeCloudSync resolves the effective mode and, if it says yes, uploads
// the PNG and inserts the row. Returns the SyncResult on success so callers
// that need the imagen.images.id (e.g. the worker linking a job row) can pick
// it up. ownerOverride, when non-empty, wins over config + env — the worker
// passes the job row's owner_user_id so each job is attributed correctly.
func maybeCloudSync(ctx context.Context, cfg *config.Config, noCloud bool, ownerOverride string, paths *output.Outputs, in output.Inputs, res *backend.Result, width, height int) (*cloud.SyncResult, error) {
	mode, err := resolveCloudSyncMode(cfg, noCloud, os.Getenv("IMAGEN_CLOUD_SYNC"))
	if err != nil {
		return nil, err
	}
	if mode == "off" {
		return nil, nil
	}

	sink, ok := cloud.NewFromEnv()
	if !ok {
		if mode == "on" {
			return nil, fmt.Errorf("cloud_sync=on but SUPABASE_URL / SUPABASE_SERVICE_KEY not set in env")
		}
		// auto + missing env = silent skip.
		return nil, nil
	}
	switch {
	case ownerOverride != "":
		sink.OwnerUserID = ownerOverride
	case cfg != nil && cfg.OwnerUserID != "":
		// Config-supplied owner_user_id takes precedence over $IMAGEN_OWNER_USER_ID.
		sink.OwnerUserID = cfg.OwnerUserID
	}
	if sink.OwnerUserID == "" {
		if mode == "on" {
			return nil, fmt.Errorf("cloud_sync=on but owner_user_id not set in config and $IMAGEN_OWNER_USER_ID is empty")
		}
		// auto + missing UUID = silent skip.
		return nil, nil
	}

	pngBytes, readErr := os.ReadFile(paths.ImagePath)
	if readErr != nil {
		return nil, fmt.Errorf("read local image: %w", readErr)
	}

	// Reuse the writer's date/slug/seed so storage_path mirrors the local
	// filename's prefix exactly — viewers can join `imagen.images` on
	// either side without timezone drift.
	date := paths.Date
	slug := paths.Slug
	if date == "" || slug == "" {
		now := time.Now()
		date = now.Format("2006-01-02")
		slug = output.Slug(in.Prompt)
	}
	ext := in.Ext
	if ext == "" {
		ext = strings.TrimPrefix(filepath.Ext(paths.ImagePath), ".")
	}
	if ext == "" {
		ext = "png"
	}

	// Snapshot the sidecar (if it exists) so the row carries the same
	// metadata view a downstream viewer would see on disk.
	var sidecar map[string]any
	if paths.SidecarPath != "" {
		if scBytes, err := os.ReadFile(paths.SidecarPath); err == nil {
			_ = json.Unmarshal(scBytes, &sidecar)
		}
	}

	model := metaString(res.Metadata, "model")
	steps := metaInt(res.Metadata, "steps")
	cost := metaFloatPtr(res.Metadata, "cost_usd_estimate")
	latency := metaInt(res.Metadata, "latency_ms")

	seed := paths.Seed
	if seed == 0 {
		seed = in.Seed
	}
	syncReq := cloud.SyncRequest{
		Date:            date,
		Slug:            slug,
		Seed:            seed,
		Ext:             ext,
		PNG:             pngBytes,
		MimeType:        res.MimeType,
		Prompt:          in.Prompt,
		Backend:         in.Backend,
		Model:           model,
		Steps:           steps,
		Width:           width,
		Height:          height,
		LatencyMs:       latency,
		CostUSDEstimate: cost,
		Sidecar:         sidecar,
	}
	syncCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()
	return sink.Sync(syncCtx, syncReq)
}

func metaString(m map[string]any, key string) string {
	if v, ok := m[key]; ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

func metaInt(m map[string]any, key string) int {
	v, ok := m[key]
	if !ok {
		return 0
	}
	switch n := v.(type) {
	case int:
		return n
	case int64:
		return int(n)
	case float64:
		return int(n)
	}
	return 0
}

func metaFloatPtr(m map[string]any, key string) *float64 {
	v, ok := m[key]
	if !ok {
		return nil
	}
	switch n := v.(type) {
	case float64:
		return &n
	case float32:
		f := float64(n)
		return &f
	case int:
		f := float64(n)
		return &f
	case int64:
		f := float64(n)
		return &f
	}
	return nil
}

// resolvePreviewMode applies the precedence chain config -> env -> flag.
// Flags win, env beats config, config beats the implicit auto default.
func resolvePreviewMode(cfg *config.Config, flagOn, flagOff bool, env string) (preview.Mode, error) {
	mode := preview.ModeAuto
	if cfg != nil && cfg.Output.Preview != "" {
		m, err := preview.ParseMode(cfg.Output.Preview)
		if err != nil {
			return "", fmt.Errorf("config output.preview: %w", err)
		}
		mode = m
	}
	if env != "" {
		m, err := preview.ParseMode(env)
		if err != nil {
			return "", fmt.Errorf("$IMAGEN_PREVIEW: %w", err)
		}
		mode = m
	}
	if flagOn && flagOff {
		return "", userErr("--preview and --no-preview are mutually exclusive")
	}
	if flagOn {
		mode = preview.ModeOn
	}
	if flagOff {
		mode = preview.ModeOff
	}
	return mode, nil
}

// maybePreview resolves the effective preview mode and, if it says yes,
// spawns a tmux window via tmux-img. Always non-fatal.
func maybePreview(cfg *config.Config, flagOn, flagOff bool, imagePath, rawPrompt string) error {
	mode, err := resolvePreviewMode(cfg, flagOn, flagOff, os.Getenv("IMAGEN_PREVIEW"))
	if err != nil {
		return err
	}
	decision, err := preview.Resolve(mode, os.Getenv("TMUX") != "", stdoutIsTTY())
	if err != nil {
		return err
	}
	if !decision.ShouldPreview {
		return nil
	}
	spawner := &preview.Spawner{}
	return spawner.Spawn(imagePath, output.Slug(rawPrompt))
}

func stdoutIsTTY() bool {
	fi, err := os.Stdout.Stat()
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeCharDevice != 0
}

// splitLeadingPositional separates the positional args at the start of args
// from the rest (which begins with the first flag). A literal "--" terminator
// pushes everything after it into the positional list and out of flag parsing.
func splitLeadingPositional(args []string) (positional, flags []string) {
	for i, a := range args {
		if a == "--" {
			return append(positional, args[i+1:]...), flags
		}
		if strings.HasPrefix(a, "-") {
			return positional, args[i:]
		}
		positional = append(positional, a)
	}
	return positional, flags
}

func parseSize(s string) (int, int, error) {
	parts := strings.SplitN(s, "x", 2)
	if len(parts) != 2 {
		return 0, 0, fmt.Errorf("expected WxH, got %q", s)
	}
	w, err := strconv.Atoi(parts[0])
	if err != nil {
		return 0, 0, err
	}
	h, err := strconv.Atoi(parts[1])
	if err != nil {
		return 0, 0, err
	}
	return w, h, nil
}

// attachUsageSink wires a Supabase cost-tracking sink into the backend
// when it accepts one and the env is configured. Adapters that record
// usage expose a public Sink field of type backend.UsageSink.
func attachUsageSink(be backend.Backend) {
	r, ok := be.(*backend.Replicate)
	if !ok {
		return
	}
	sink, ok := usage.NewSupabaseSinkFromEnv()
	if !ok {
		return
	}
	r.Sink = sink
}

func buildBackend(cfg *config.Config, name string) (backend.Backend, error) {
	if cfg != nil {
		spec, ok := cfg.Backends[name]
		if ok {
			return backend.Default.Build(spec.Type, name, spec.Raw)
		}
	}
	if backend.Default.Has(name) {
		return backend.Default.Build(name, name, nil)
	}
	return nil, userErr("backend %q not found in config and not a registered type (registered types: %v)",
		name, backend.Default.Types())
}

func buildWriter(cfg *config.Config, noSidecar bool) *output.Writer {
	w := &output.Writer{}
	if cfg != nil {
		w.Directory = config.ExpandPath(cfg.Output.Directory)
		w.NameTemplate = cfg.Output.Naming
		w.WriteSidecar = cfg.Output.WriteMetadataJSON
	}
	if w.Directory == "" {
		w.Directory = "."
	}
	if noSidecar {
		w.WriteSidecar = false
	}
	return w
}

func seedFromMetadata(meta map[string]any, fallback int64) int64 {
	if v, ok := meta["seed"]; ok {
		switch n := v.(type) {
		case int64:
			return n
		case int:
			return int64(n)
		case float64:
			return int64(n)
		}
	}
	return fallback
}

func extFromMime(mime string) string {
	switch mime {
	case "image/png", "":
		return "png"
	case "image/jpeg":
		return "jpg"
	case "image/webp":
		return "webp"
	}
	return "bin"
}