package errgroups

import (
	"context"
	"dashboard/utils"
	"sync"

	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

type ErrGroupFunc struct {
	wg   *errgroup.Group
	mu   sync.RWMutex
	funs map[string]FuncWithErrFunc
}

type FuncWithErrFunc func(context.Context) error

// 注意使用func时，如果想要对数组对象执行多个函数，不可以使用此种方法，因为函数名都是一样的，map里面会被替换。此种场景需要使用接口的方式。
func NewErrGroupFunc() *ErrGroupFunc {
	return &ErrGroupFunc{
		funs: make(map[string]FuncWithErrFunc),
	}
}

func (e *ErrGroupFunc) Run(ctx context.Context) error {
	e.mu.RLock()
	defer e.mu.RUnlock()

	e.wg, ctx = errgroup.WithContext(ctx)

	for name, fn := range e.funs {
		fn := fn
		name := name
		e.wg.Go(func() error {
			defer zap.L().Sugar().Warnf("func: %s stop", name)
			zap.L().Sugar().Infof("func: %s runing...", name)
			return fn(ctx)
		})
	}

	if err := e.wg.Wait(); err != nil {
		return err
	}

	return nil
}

func (e *ErrGroupFunc) Add(fn ...FuncWithErrFunc) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for _, f := range fn {
		fnStr := utils.GetFuncName(f)
		e.funs[fnStr] = f
	}
}

func (e *ErrGroupFunc) Del(fn ...FuncWithErrFunc) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for _, f := range fn {
		fnStr := utils.GetFuncName(f)
		delete(e.funs, fnStr)
	}
}

type ErrGroup struct {
	wg   *errgroup.Group
	mu   sync.RWMutex
	funs map[FuncWithErr]struct{}
}

type FuncWithErr interface {
	Run(context.Context) error
}

func NewErrGroup() *ErrGroup {
	return &ErrGroup{
		funs: make(map[FuncWithErr]struct{}),
	}
}

func (e *ErrGroup) Run(ctx context.Context) error {
	e.mu.RLock()
	defer e.mu.RUnlock()

	e.wg, ctx = errgroup.WithContext(ctx)

	for fn := range e.funs {
		fn := fn // shadow
		name := utils.GetInterfaceName(fn)
		e.wg.Go(func() error {
			defer zap.L().Sugar().Warnf("Object: %s stop", name)
			zap.L().Sugar().Infof("Object: %s runing...", name)
			return fn.Run(ctx)
		})
	}

	if err := e.wg.Wait(); err != nil {
		return err
	}

	return nil
}

func (e *ErrGroup) Add(fn ...FuncWithErr) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for _, f := range fn {
		e.funs[f] = struct{}{}
	}
}

func (e *ErrGroup) Del(fn ...FuncWithErr) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for _, f := range fn {
		delete(e.funs, f)
	}
}