container.go

package pal

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"maps"
	"reflect"
	"slices"

	typetostring "github.com/samber/go-type-to-string"

	"golang.org/x/sync/errgroup"

	"github.com/zhulik/pal/pkg/dag"
)

type factoryServiceMaping struct {
	Factory any
	Service ServiceDef
}

type factoryService interface {
	Factory() any
	MustFactory() any
}

// Container is responsible for storing services, instances and the dependency graph
type Container struct {
	pal *Pal

	services  map[string]ServiceDef
	factories map[string]factoryServiceMaping
	graph     *dag.DAG[string, ServiceDef]
	logger    *slog.Logger
}

// NewContainer creates a new Container instance
func NewContainer(pal *Pal, services ...ServiceDef) *Container {
	services = flattenServices(services)

	container := &Container{
		pal:       pal,
		services:  map[string]ServiceDef{},
		factories: map[string]factoryServiceMaping{},
		graph:     dag.New[string, ServiceDef](),
		logger:    slog.With("palComponent", "Container"),
	}

	for _, service := range services {
		container.addService(service)
		if factory, ok := service.(factoryService); ok {
			// Add Factory to the container
			fn := factory.Factory()
			fnType := reflect.TypeOf(fn)
			container.factories[typetostring.GetReflectType(fnType)] = factoryServiceMaping{
				Factory: fn,
				Service: service,
			}

			// Add MustFactory to the container
			fn = factory.MustFactory()
			fnType = reflect.TypeOf(fn)
			container.factories[typetostring.GetReflectType(fnType)] = factoryServiceMaping{
				Factory: fn,
				Service: service,
			}
		}
	}

	return container
}

func (c *Container) Init(ctx context.Context) error {
	c.logger.Debug("Building dependency tree...")

	for _, service := range c.services {
		if err := c.addDependencyVertex(service, nil); err != nil {
			return err
		}
	}

	for _, service := range c.graph.ReverseTopologicalOrder() {
		if err := service.Init(ctx); err != nil {
			c.logger.Error("Failed to initialize container", "error", err)
			return err
		}
	}

	c.logger.Debug("Container initialized")
	return nil
}

func (c *Container) Invoke(ctx context.Context, name string, args ...any) (any, error) {
	service, ok := c.services[name]
	if !ok {
		return nil, fmt.Errorf("%w: '%s', known services: %s", ErrServiceNotFound, name, c.services)
	}

	if len(args) != service.Arguments() {
		return nil, fmt.Errorf("%w: '%s': %d arguments expected, got %d", ErrServiceInvalidArgumentsCount, name, service.Arguments(), len(args))
	}

	instance, err := service.Instance(ctx, args...)
	if err != nil {
		return nil, fmt.Errorf("%w: '%s': %w", ErrServiceInitFailed, name, err)
	}

	return instance, nil
}

func (c *Container) InvokeByInterface(ctx context.Context, iface reflect.Type, args ...any) (any, error) {
	if iface.Kind() != reflect.Interface {
		return nil, fmt.Errorf("%w: must be an interface, got %s", ErrNotAnInterface, iface.String())
	}

	matches := []ServiceDef{}
	for _, service := range c.services {
		instance := service.Make()
		if instance == nil {
			continue
		}
		if reflect.TypeOf(instance).Implements(iface) {
			matches = append(matches, service)
		}
	}
	if len(matches) == 0 {
		return nil, fmt.Errorf("%w: no implementations of %s found", ErrServiceNotFound, iface.String())
	}

	if len(matches) == 1 {
		return matches[0].Instance(ctx, args...)
	}

	return nil, fmt.Errorf("%w: found %d services for interface %s", ErrMultipleServicesFoundByInterface, len(matches), iface.String())
}

func (c *Container) InjectInto(ctx context.Context, target any) error {
	v := reflect.ValueOf(target).Elem()
	t := v.Type()

	for i := 0; i < t.NumField(); i++ {
		field := v.Field(i)

		tags, err := ParseTag(t.Field(i).Tag.Get("pal"))
		if err != nil {
			return err
		}
		if _, ok := tags[TagSkip]; ok || !field.CanSet() {
			continue
		}

		fieldType := t.Field(i).Type

		if fieldType == reflect.TypeOf((*slog.Logger)(nil)) && c.pal.config.AttrSetters != nil {
			c.injectLoggerIntoField(field, target)
			continue
		}

		if _, ok := tags[TagMatchInterface]; ok {
			err = c.injectByInterface(ctx, field, fieldType)
			if err != nil {
				return err
			}
			continue
		}

		typeName, mustInject := tags[TagName]

		if typeName == "" {
			typeName = typetostring.GetReflectType(fieldType)
		}

		if fieldType.Kind() == reflect.Func {
			mapping, ok := c.factories[typeName]
			if ok {
				field.Set(reflect.ValueOf(mapping.Factory))
			}

			continue
		}

		err = c.injectByName(ctx, typeName, field)
		if err != nil {
			if errors.Is(err, ErrServiceNotFound) && !mustInject {
				continue
			}
			return err
		}
	}

	return nil
}

func (c *Container) injectByInterface(ctx context.Context, field reflect.Value, fieldType reflect.Type) error {
	dependency, err := c.InvokeByInterface(ctx, fieldType)
	if err != nil {
		return err
	}

	field.Set(reflect.ValueOf(dependency))

	return nil
}

func (c *Container) injectByName(ctx context.Context, name string, field reflect.Value) error {
	dependency, err := c.Invoke(ctx, name)
	if err != nil {
		if errors.Is(err, ErrServiceInvalidArgumentsCount) {
			return fmt.Errorf("%w: '%s': %w", ErrFactoryServiceDependency, name, err)
		}
		return err
	}

	field.Set(reflect.ValueOf(dependency))

	return nil
}

func (c *Container) Shutdown(ctx context.Context) error {
	c.logger.Debug("Shutting down all runners")

	for _, service := range c.graph.TopologicalOrder() {
		err := service.Shutdown(ctx)
		if err != nil {
			c.logger.Error("Failed to shutdown service. Exiting immediately", "service", service.Name(), "error", err)
			return err
		}
	}

	c.logger.Debug("Container shut down successfully")
	return nil
}

func (c *Container) HealthCheck(ctx context.Context) error {
	var wg errgroup.Group

	c.logger.Debug("Healthchecking services")

	for _, service := range c.graph.TopologicalOrder() {
		wg.Go(func() error {
			// Do not check pal again, this leads to recursion
			if service.Name() == "*github.com/zhulik/pal.Pal" {
				return nil
			}

			err := service.HealthCheck(ctx)
			if err != nil {
				return err
			}

			return nil
		})
	}

	err := wg.Wait()
	if err != nil {
		c.logger.Error("Healthcheck failed", "error", err)
		return err
	}

	c.logger.Debug("Healthcheck successful")

	return nil
}

// Services returns a map of all registered services in the container, keyed by their names.
// This can be useful for debugging or introspection purposes.
func (c *Container) Services() map[string]ServiceDef {
	return c.services
}

// StartRunners starts all services that implement the Runner interface in background goroutines.
// It creates a cancellable context that will be canceled during shutdown.
// Returns an error if any runner fails, though runners continue to execute independently.
func (c *Container) StartRunners(ctx context.Context) error {
	services := slices.Collect(maps.Values(c.services))
	return RunServices(ctx, services)
}

// Graph returns the dependency graph of services.
// This can be useful for visualization or analysis of the service dependencies.
func (c *Container) Graph() *dag.DAG[string, ServiceDef] {
	return c.graph
}

func (c *Container) addService(service ServiceDef) {
	setPalField(reflect.ValueOf(service), c.pal, map[reflect.Value]bool{})
	c.services[service.Name()] = service
}

// addDependencyVertex adds a service to the dependency graph and recursively adds its dependencies.
// If parent is not nil, it also adds an edge from parent to service in the graph.
// This method is used during container initialization to build the complete dependency graph.
func (c *Container) addDependencyVertex(service ServiceDef, parent ServiceDef) error {
	c.graph.AddVertexIfNotExist(service.Name(), service)

	if parent != nil {
		if err := c.graph.AddEdgeIfNotExist(parent.Name(), service.Name()); err != nil {
			return err
		}
	}
	m := service.Make()
	if isNil(m) {
		return nil
	}
	val := reflect.ValueOf(m)
	if val.Kind() == reflect.Ptr {
		val = val.Elem()
	}

	if !val.IsValid() {
		return nil
	}

	typ := val.Type()
	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)

		if !field.IsExported() {
			continue
		}

		tags, err := ParseTag(field.Tag.Get("pal"))
		if err != nil {
			return err
		}

		dependencyName := tags[TagName]

		if dependencyName == "" {
			dependencyName = typetostring.GetReflectType(field.Type)
		}

		if childService, ok := c.services[dependencyName]; ok {
			if err := c.addDependencyVertex(childService, service); err != nil {
				return err
			}
		}

		if factoryMapping, ok := c.factories[dependencyName]; ok {
			if err := c.addDependencyVertex(factoryMapping.Service, service); err != nil {
				return err
			}
		}
	}

	return nil
}

func (c *Container) injectLoggerIntoField(field reflect.Value, target any) {
	logger := slog.Default()
	for _, attrSetter := range c.pal.config.AttrSetters {
		name, value := attrSetter(target)
		logger = logger.With(name, value)
	}
	field.Set(reflect.ValueOf(logger))
}

func setPalField(v reflect.Value, pal *Pal, visited map[reflect.Value]bool) {
	if visited[v] {
		return
	}
	visited[v] = true

	if v.Kind() == reflect.Interface {
		v = v.Elem()
	}

	if v.Kind() == reflect.Pointer {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < v.NumField(); i++ {
		field := v.Field(i)
		if field.CanSet() && field.Type() == reflect.TypeOf(pal) {
			field.Set(reflect.ValueOf(pal))
		}

		if field.Kind() == reflect.Struct || (field.Kind() == reflect.Pointer && !field.IsNil()) {
			setPalField(field, pal, visited)
		}

		if field.Kind() == reflect.Array || field.Kind() == reflect.Slice {
			for i := 0; i < field.Len(); i++ {
				item := field.Index(i)
				setPalField(item, pal, visited)
			}
		}
	}
}