package graceful_shutdown

import (
	"context"
	"git.rinsvent.ru/rinsvent/gol"
	"go.uber.org/zap"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

type ContextShutdowner struct {
	ctx    context.Context
	cancel context.CancelFunc
	wg     *sync.WaitGroup
}

func (sh *ContextShutdowner) IncLock() {
	sh.wg.Add(1)
}

func (sh *ContextShutdowner) DecLock() {
	sh.wg.Add(-1)
}

func (sh *ContextShutdowner) GetContext() context.Context {
	return sh.ctx
}

func (sh *ContextShutdowner) Wait() {
	sigChan := make(chan os.Signal, 2)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	sig := <-sigChan

	gol.With(zap.String("sig", sig.String())).Info("Gracefully shutting down")

	sh.cancel()
	sh.wg.Wait()
}

func NewContextShutdowner() ContextShutdowner {
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	return ContextShutdowner{
		ctx:    ctx,
		cancel: cancel,
		wg:     &wg,
	}
}