package goshutdowner

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

// Simple struct, use like a sync.Mutex
//
//  var s goshutdowner.Shutdowner
//  s.Add(func() {
//	 log.Println("Shutting down")
//  })
type Shutdowner struct {
	initialized bool
	quit        chan os.Signal
	funcs       []ShutdownFunc
	wg          sync.WaitGroup
	mutex       sync.RWMutex
}

type ShutdownFunc func()

// Internal method
func (s *Shutdowner) init() {
	if s.initialized {
		return
	}
	s.quit = make(chan os.Signal, 1)
	signal.Notify(s.quit,
		os.Interrupt,
		syscall.SIGINT,
		syscall.SIGTERM, // e.g. Docker stop
	)
	go func() {
		<-s.quit
		s.Shutdown()
	}()
	s.initialized = true
}

// Add a func, that should be called with s.Shutdown() or when receiving a shutdown signal
func (s *Shutdowner) Add(f ShutdownFunc) {
	s.init()
	s.mutex.Lock()
	s.wg.Add(1)
	s.funcs = append(s.funcs, f)
	s.mutex.Unlock()
}

// Trigger shutdown directly
func (s *Shutdowner) Shutdown() {
	s.init()
	s.mutex.RLock()
	for _, f := range s.funcs {
		go func(f func()) {
			defer s.wg.Done()
			f()
		}(f)
	}
	s.mutex.RUnlock()
	s.wg.Wait()
}

// Wait till all functions finished
func (s *Shutdowner) Wait() {
	s.init()
	s.wg.Wait()
}

// Shutdown and wait till shutdown finished. Shorthand for:
//
//  s.Shutdown()
//  s.Wait()
func (s *Shutdowner) ShutdownAndWait() {
	s.Shutdown()
	s.Wait()
}