package syncz

import (
	"context"
	"log/slog"

	"github.com/ash2k/stager"
)

// Layer is a layer in a startup and shutdown stack.
// Each layer can expand into zero or more stages that are run first to last and shut down in reverse order.
//
// # Example
//
// |---------|
// | Layer A |
// |---------|
// | Layer B |
// |---------|
//
// Let's say the above layers expand into the following:
//
// |----------|
// | Stage A1 |
// |----------|
// | Stage A2 |
// |----------|
// | Stage B1 |
// |----------|
// | Stage B2 |
// |----------|
//
// Stages are shut down in bottom up order: B2, B1, A2, A1.
//
// While stages are started in the order they come, it doesn't matter much since they all start concurrently in
// independent goroutines and, hence, no real ordering is achieved.
// The purpose of this mechanism is to guarantee reverse shutdown order of the stages.
// This is ensured by the underlying github.com/ash2k/stager library.
// Layers provide more flexibility and composability on top of the library.
type Layer interface {
	ToStageFuncs() ([]stager.StageFunc, error)
}

var (
	_ Layer = LayerFunc(nil)
	_ Layer = LayerGoFunc(nil)
	_ Layer = SiblingLayers(nil)
	_ Layer = (*LogWhenDoneLayer)(nil)
)

func RunLayers(ctx context.Context, layers ...Layer) error {
	sfs, err := layers2funcs(layers)
	if err != nil {
		return err
	}
	return stager.RunStages(ctx, sfs...)
}

type LayerFunc stager.StageFunc

func (f LayerFunc) ToStageFuncs() ([]stager.StageFunc, error) {
	return []stager.StageFunc{
		stager.StageFunc(f),
	}, nil
}

type LayerGoFunc func(ctx context.Context) error

func (f LayerGoFunc) ToStageFuncs() ([]stager.StageFunc, error) {
	return []stager.StageFunc{
		func(stage stager.Stage) {
			stage.Go(f)
		},
	}, nil
}

type LogWhenDoneLayer struct {
	Log *slog.Logger
	Msg string
}

func (l *LogWhenDoneLayer) ToStageFuncs() ([]stager.StageFunc, error) {
	return []stager.StageFunc{
		func(stage stager.Stage) {
			stage.GoWhenDone(func() error {
				l.Log.Info(l.Msg)
				return nil
			})
		},
	}, nil
}

// SiblingLayers allows to compose multiple layers to be siblings rather than ordered on top of each other.
//
// # Example
//
// A is a simple layer and B and C are sibling layers.
// |---------|---------|
// | Layer A           |
// |---------|---------|
// | Layer B | Layer C |
// |---------|---------|
//
// Can expand into:
//
// |----------|----------|
// | Stage A1            |
// |----------|----------|
// | Stage B1 | Stage C1 |
// |          |----------|
// |          | Stage C2 |
// |----------|----------|
//
// Shutdown order: (B1 and (C2 then C1)) then A1. Where "and" means things happen concurrently and "then" means
// an ordering dependency.
type SiblingLayers []Layer

func (l SiblingLayers) ToStageFuncs() ([]stager.StageFunc, error) {
	var sfs []stager.StageFunc
	for _, layer := range l {
		sf, err := layer.ToStageFuncs()
		if err != nil {
			return nil, err
		}
		switch len(sf) {
		case 0:
			// Nothing to run
		case 1:
			sfs = append(sfs, sf[0])
		default:
			sfs = append(sfs, func(stage stager.Stage) {
				// Ensure all the returned staged funcs run in the correct order but use the same single stage
				// as the other layers.
				stage.Go(func(ctx context.Context) error {
					return stager.RunStages(ctx, sf...)
				})
			})
		}
	}
	return []stager.StageFunc{
		func(stage stager.Stage) { // single stage that runs all layers
			for _, sf := range sfs {
				sf(stage)
			}
		},
	}, nil
}

func layers2funcs(layers []Layer) ([]stager.StageFunc, error) {
	var sfs []stager.StageFunc
	for _, layer := range layers {
		sf, err := layer.ToStageFuncs()
		if err != nil {
			return nil, err
		}
		sfs = append(sfs, sf...)
	}
	return sfs, nil
}
