ImaGen/internal/output/output.go

// Package output writes generated images to disk and (optionally) a JSON
// metadata sidecar. Filenames are resolved through a small template language
// with placeholders {date}, {time}, {slug}, {seed}, {backend}, {ext}.
package output

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

// Writer renders one generation to disk under Directory.
type Writer struct {
	Directory     string
	NameTemplate  string
	WriteSidecar  bool
	Now           func() time.Time
}

// Inputs are the ingredients needed to compute a filename and write a sidecar.
type Inputs struct {
	Prompt   string
	Backend  string
	Seed     int64
	Ext      string
	Metadata map[string]any
}

// Outputs lists the artefacts the writer produced.
type Outputs struct {
	ImagePath    string
	SidecarPath  string
	// Date is the YYYY-MM-DD the writer used for the filename. Cloud sync
	// reuses this so storage_path matches the local filename's date.
	Date string
	// Slug is the filename-safe prompt fragment the writer used.
	Slug string
	// Seed is the seed value baked into the filename.
	Seed int64
}

// Write streams img to disk and, if enabled, writes a sidecar. The image
// stream is consumed even on error so callers don't leak goroutines from
// piped readers.
func (w *Writer) Write(img io.Reader, in Inputs) (*Outputs, error) {
	now := w.now()
	ext := in.Ext
	if ext == "" {
		ext = "png"
	}
	tmpl := w.NameTemplate
	if tmpl == "" {
		tmpl = "{date}-{slug}-{seed}.{ext}"
	}
	date := now.Format("2006-01-02")
	slug := Slug(in.Prompt)
	name := renderTemplate(tmpl, map[string]string{
		"date":    date,
		"time":    now.Format("150405"),
		"slug":    slug,
		"seed":    fmt.Sprintf("%d", in.Seed),
		"backend": in.Backend,
		"ext":     strings.TrimPrefix(ext, "."),
	})

	dir := w.Directory
	if dir == "" {
		dir = "."
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("mkdir %s: %w", dir, err)
	}

	imagePath := filepath.Join(dir, name)
	f, err := os.Create(imagePath)
	if err != nil {
		return nil, fmt.Errorf("create %s: %w", imagePath, err)
	}
	if _, err := io.Copy(f, img); err != nil {
		f.Close()
		return nil, fmt.Errorf("write %s: %w", imagePath, err)
	}
	if err := f.Close(); err != nil {
		return nil, fmt.Errorf("close %s: %w", imagePath, err)
	}

	out := &Outputs{ImagePath: imagePath, Date: date, Slug: slug, Seed: in.Seed}

	if w.WriteSidecar {
		sidecar := imagePath + ".json"
		body := map[string]any{
			"timestamp": now.UTC().Format(time.RFC3339),
			"prompt":    in.Prompt,
			"backend":   in.Backend,
			"seed":      in.Seed,
			"image":     filepath.Base(imagePath),
			"metadata":  in.Metadata,
		}
		data, err := json.MarshalIndent(body, "", "  ")
		if err != nil {
			return out, fmt.Errorf("marshal sidecar: %w", err)
		}
		if err := os.WriteFile(sidecar, append(data, '\n'), 0o644); err != nil {
			return out, fmt.Errorf("write sidecar %s: %w", sidecar, err)
		}
		out.SidecarPath = sidecar
	}
	return out, nil
}

// WriteToPath bypasses templating and writes img to an explicit path. This is
// the path the CLI's --output flag uses.
func (w *Writer) WriteToPath(img io.Reader, path string, in Inputs) (*Outputs, error) {
	now := w.now()
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, fmt.Errorf("mkdir %s: %w", filepath.Dir(path), err)
	}
	f, err := os.Create(path)
	if err != nil {
		return nil, fmt.Errorf("create %s: %w", path, err)
	}
	if _, err := io.Copy(f, img); err != nil {
		f.Close()
		return nil, fmt.Errorf("write %s: %w", path, err)
	}
	if err := f.Close(); err != nil {
		return nil, fmt.Errorf("close %s: %w", path, err)
	}
	out := &Outputs{ImagePath: path, Date: now.Format("2006-01-02"), Slug: Slug(in.Prompt), Seed: in.Seed}
	if w.WriteSidecar {
		sidecar := path + ".json"
		body := map[string]any{
			"timestamp": now.UTC().Format(time.RFC3339),
			"prompt":    in.Prompt,
			"backend":   in.Backend,
			"seed":      in.Seed,
			"image":     filepath.Base(path),
			"metadata":  in.Metadata,
		}
		data, err := json.MarshalIndent(body, "", "  ")
		if err != nil {
			return out, fmt.Errorf("marshal sidecar: %w", err)
		}
		if err := os.WriteFile(sidecar, append(data, '\n'), 0o644); err != nil {
			return out, fmt.Errorf("write sidecar %s: %w", sidecar, err)
		}
		out.SidecarPath = sidecar
	}
	return out, nil
}

func (w *Writer) now() time.Time {
	if w.Now != nil {
		return w.Now()
	}
	return time.Now()
}

var (
	tmplPlaceholder = regexp.MustCompile(`\{([a-z]+)\}`)
	slugAllowed     = regexp.MustCompile(`[^a-z0-9]+`)
)

func renderTemplate(t string, vars map[string]string) string {
	return tmplPlaceholder.ReplaceAllStringFunc(t, func(match string) string {
		key := match[1 : len(match)-1]
		if v, ok := vars[key]; ok {
			return v
		}
		return match
	})
}

// Slug normalises a prompt fragment into a filesystem-safe token.
func Slug(s string) string {
	s = strings.ToLower(s)
	s = slugAllowed.ReplaceAllString(s, "-")
	s = strings.Trim(s, "-")
	if s == "" {
		s = "image"
	}
	const max = 40
	if len(s) > max {
		s = s[:max]
		s = strings.TrimRight(s, "-")
	}
	return s
}