package taskmanager

import (
	"reflect"
	"sync"

	log "github.com/Sirupsen/logrus"
)

// BackgroundTask is the interface a task must implement to be managed.
type BackgroundTask interface {
	Start()
	Stop()
	Done() <-chan error
}

// TaskManager starts up tasks, waits for them to exit, shuts them all down if any exit.
type TaskManager struct {
	tasks   []BackgroundTask
	wg      sync.WaitGroup
	started bool
}

// New creates an initialized Taskmanager.
func New() *TaskManager {
	return &TaskManager{tasks: make([]BackgroundTask, 0)}
}

// AddTask adds a task to the manager.
// Don't use this after the manager is started.
func (m *TaskManager) AddTask(t BackgroundTask) {
	if m.started {
		log.Fatal("Can't add a task after the manager is started")
	}
	m.tasks = append(m.tasks, t)
	m.wg.Add(1)
}

// StartTasks calls .Start() on each task and then starts watching them for exits.
func (m *TaskManager) StartTasks() {
	if m.started {
		return
	}
	m.started = true
	for _, t := range m.tasks {
		t.Start()
	}
	go m.watchTasks()
}

func (m *TaskManager) watchTasks() {
	cases := make([]reflect.SelectCase, len(m.tasks))
	for i, t := range m.tasks {
		cases[i] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(t.Done())}
	}
	remaining := len(cases)
	for remaining > 0 {
		chosen, _, ok := reflect.Select(cases)
		if !ok {
			// The channel has been closed, so zero out the channel to disable the case.
			cases[chosen].Chan = reflect.ValueOf(nil)
			remaining--
			m.wg.Done()
		}
		m.StopAll()
	}
}

// Wait is called after StartTasks to wait for the tasks to all exit.
func (m *TaskManager) Wait() {
	m.wg.Wait()
}

// StopAll calls .Stop() on each task.
// If any task exists, the manager will do this.
func (m *TaskManager) StopAll() {
	for _, t := range m.tasks {
		t.Stop()
	}
}
