app/lifecycle.go

package app

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sort"
	"strings"
)

type contextKey string

const lifecycleContextKey contextKey = "lifecycle"

// LifecycleOpts contains user-exposed options when defining a lifecycle.
type LifecycleOpts struct {
	// DisableAutoload disables dependency autoloading (enabled by default).
	DisableAutoload bool

	// Logger is the logger for the lifecycle. If not set, it will be initialized
	// by the lifecycle using the default logger (slog.Default).
	Logger *slog.Logger
}

// Lifecycle represents the core application structure. Lifecycle manages
// resources providing an orchestrator for setup and teardown of modules.
type Lifecycle struct {
	modules []*Module
	opts    LifecycleOpts

	setupCount      int
	setupTracker    map[string]int
	teardownCount   int
	teardownTracker map[string]int
}

// NewLifecycle creates a new Lifecycle instance with a default logger and the
// given modules. It panics if any module has a duplicate name.
func NewLifecycle(modules ...*Module) *Lifecycle {
	// Ensure modules are unique
	unique := make(map[string]bool)
	for _, mod := range modules {
		if _, exists := unique[mod.name]; exists {
			panic(fmt.Sprintf("duplicate module: %s", mod))
		}
		unique[mod.name] = true
	}

	return &Lifecycle{
		modules:         modules,
		setupTracker:    make(map[string]int),
		teardownTracker: make(map[string]int),
	}
}

// LifecycleToContext adds the Lifecycle to the context.
func LifecycleToContext(ctx context.Context, lifecycle *Lifecycle) context.Context {
	if lifecycle == nil {
		return ctx
	}
	return context.WithValue(ctx, lifecycleContextKey, lifecycle)
}

// LifecycleFromContext retrieves the Lifecycle from the context. Returns nil if not found.
func LifecycleFromContext(ctx context.Context) *Lifecycle {
	if lifecycle, ok := ctx.Value(lifecycleContextKey).(*Lifecycle); ok {
		return lifecycle
	}
	return nil
}

// WithOpts sets the options for the lifecycle.
func (app *Lifecycle) WithOpts(opts LifecycleOpts) *Lifecycle {
	app.opts = opts
	return app
}

// WithLogger sets the logger for the lifecycle. Panics if the logger is nil.
func (app *Lifecycle) WithLogger(logger *slog.Logger) *Lifecycle {
	if logger == nil {
		panic("logger cannot be nil")
	}
	app.opts.Logger = logger
	return app
}

// Logger returns the logger for the lifecycle.
func (app *Lifecycle) Logger() *slog.Logger {
	if app == nil {
		panic("lifecycle is nil, cannot get logger")
	}
	if app.opts.Logger == nil {
		app.opts.Logger = slog.Default()
	}
	return app.opts.Logger
}

// Setup initializes all modules in the order they were defined and checks
// for dependencies. It returns an error if any module fails to initialize or
// if dependencies are not satisfied. Lifecycle.Teardown should always be run at the end
// of the application lifecycle to ensure all resources are cleaned up properly.
func (app *Lifecycle) Setup() error {
	for _, mod := range app.modules {
		if err := app.setupSingle(nil, mod); err != nil {
			return err
		}
	}

	app.Logger().Info("Lifecycle modules initialized, backend setup complete",
		"modules", mapToString(app.setupTracker))

	return nil
}

// Teardown runs all module teardown functions in reverse order of setup.
// Teardown should always be run at the end of the application lifecycle
// to ensure all resources are cleaned up properly. All module tear down
// errors are returned as a single error (non-blocking).
func (app *Lifecycle) Teardown() error {
	var err error
	for i := len(app.modules) - 1; i >= 0; i-- {
		if singleErr := app.teardownSingle(app.modules[i]); singleErr != nil {
			err = errors.Join(err, singleErr)
		}
	}

	if err != nil {
		app.Logger().Error("Error tearing down modules", "failures", app.setupCount-app.teardownCount, "error", err)
		return err
	}

	app.Logger().Info("All modules torn down, backend teardown complete",
		"modules", mapToString(app.teardownTracker))

	return nil
}

// Require adds module(s) to the lifecycle and immediately runs any setup
// functions. Relevant when a module is not part of the main application
// but may still be conditionally necessary. Any modules that are already
// set up are ignored.
func (app *Lifecycle) Require(modules ...*Module) error {
	return app.require(nil, false, modules...)
}

// RequireUnique is the same as Require, but it returns an error if any requested
// module is already set up--rather than ignoring it.
func (app *Lifecycle) RequireUnique(modules ...*Module) error {
	return app.require(nil, true, modules...)
}

// RequireL adds module(s) to the lifecycle with a specific logger and
// immediately runs any setup functions. See Require for more details.
// This variation is useful when you need to set up modules with a non-
// default logger.
func (app *Lifecycle) RequireL(logger *slog.Logger, modules ...*Module) error {
	if logger == nil {
		return fmt.Errorf("logger cannot be nil")
	}
	return app.require(logger, false, modules...)
}

// RequireUniqueL is the same as RequireL, but it returns an error if any requested
// module is already set up--rather than ignoring it.
func (app *Lifecycle) RequireUniqueL(logger *slog.Logger, modules ...*Module) error {
	if logger == nil {
		return fmt.Errorf("logger cannot be nil")
	}
	return app.require(logger, true, modules...)
}

// require is a helper function that attempts to add module(s) to the lifecycle
// and immediately run any setup functions.
func (app *Lifecycle) require(logger *slog.Logger, unique bool, modules ...*Module) error {
	if len(modules) == 0 {
		return fmt.Errorf("no modules to require")
	}

	for i, mod := range modules {
		if mod == nil {
			return fmt.Errorf("module %d is nil", i)
		}

		// Check if the module has already been set up
		if _, ok := app.setupTracker[mod.name]; ok {
			if unique {
				return fmt.Errorf("module %s is already set up, cannot require again", mod)
			}

			app.Logger().Warn("module already set up, ignoring", "module", mod)

			// Mark duplicate module as loaded
			mod.loaded = true
			mod.lifecycle = app
			mod.logger = logger

			continue
		}

		// Add the module to the lifecycle
		app.modules = append(app.modules, mod)

		// Run the setup function for the module
		if err := app.setupSingle(logger, mod); err != nil {
			return fmt.Errorf("error setting up required module %s: %w", mod, err)
		}
	}

	app.Logger().Info("New modules initialized", "all", mapToString(app.setupTracker))
	return nil
}

// setupSingle is a helper function to set up a single module. Returns an error
// if the module cannot be set up or if dependencies are not satisfied.
func (app *Lifecycle) setupSingle(logger *slog.Logger, mod *Module) error {
	if mod == nil {
		return fmt.Errorf("module is nil")
	}

	// Set the parent lifecycle and logger override
	mod.lifecycle = app
	mod.logger = logger

	// Check if all dependencies are satisfied
	for _, dep := range mod.depends {
		if _, ok := app.setupTracker[dep]; !ok {
			if app.opts.DisableAutoload {
				return fmt.Errorf("dependency %s not satisfied for '%s'", dep, mod)
			} else {
				// Attempt to set up the dependency
				depmod, err := app.getModuleByName(dep)
				if err != nil {
					return fmt.Errorf("error getting dependency '%s' for %s: %w", dep, mod, err)
				}

				if err := app.setupSingle(logger, depmod); err != nil {
					return fmt.Errorf("error setting up dependency %s for %s: %w", depmod, mod, err)
				}
			}
		}
	}

	if mod.setup != nil {
		// Run the setup function for the module
		if err := mod.setup(mod); err != nil {
			return fmt.Errorf("error initializing %s: %w", mod, err)
		}
	}

	// Mark this module as setup
	app.setupTracker[mod.name] = app.setupCount
	app.setupCount++
	mod.loaded = true
	return nil
}

// teardownSingle is a helper function to tear down a single module. Returns
// an error if anything goes wrong.
func (app *Lifecycle) teardownSingle(mod *Module) error {
	if mod == nil {
		return fmt.Errorf("module is nil")
	}

	// Check if the module was set up
	if _, ok := app.setupTracker[mod.name]; !ok {
		return fmt.Errorf("module %s is not set up, cannot tear down", mod)
	}

	// Run the teardown function for the module
	if mod.teardown != nil {
		if err := mod.teardown(mod); err != nil {
			return fmt.Errorf("error tearing down %s: %w", mod, err)
		}
	}

	// Mark this module as torn down
	app.teardownTracker[mod.name] = app.teardownCount
	app.teardownCount++
	mod.loaded = false
	return nil
}

// getModuleByName retrieves a module by its name from the lifecycle.
func (app *Lifecycle) getModuleByName(name string) (*Module, error) {
	for _, mod := range app.modules {
		if mod.name == name {
			return mod, nil
		}
	}
	return nil, ErrModuleNotFound
}

// mapToString converts a map to an ordered, opinionated string representation.
func mapToString(m map[string]int) string {
	if len(m) == 0 {
		return "[]"
	}

	// Sort the map by value
	values := make([]int, 0, len(m))
	reverseMap := make(map[int]string, len(m))
	for k, v := range m {
		values = append(values, v)
		reverseMap[v] = k
	}
	sort.Ints(values)

	result := make([]string, 0, len(m))
	for _, v := range values {
		if name, ok := reverseMap[v]; ok {
			result = append(result, fmt.Sprintf("'%s'", name))
		}
	}

	return "[" + strings.Join(result, " ") + "]"
}