package grace

import (
	"context"
	"os"
	"os/signal"
	"sync"
)

type Grace struct {
	process context.Context
	cancel  func()
	wg      sync.WaitGroup
}

func New() *Grace {
	c, cancelFunc := context.WithCancel(context.Background())
	g := &Grace{
		process: c,
		cancel:  cancelFunc,
	}
	return g
}

func (g *Grace) Notify(signals ...os.Signal) {
	g.process, g.cancel = signal.NotifyContext(g.process, signals...)
}

func (g *Grace) Start(f func(context.Context)) {
	g.wg.Add(1)
	go func() {
		defer g.wg.Done()
		f(g.process)
	}()
}

func (g *Grace) StopAndWait() {
	g.cancel()
	g.Wait()
}

func (g *Grace) Wait() {
	g.wg.Wait()
}