package executor

import (
	"fmt"
	"sync"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type Task interface {
	Run()
}

// Executor defines the actions associated with the Goroutine pool.
type Executor interface {
	// Execute method corresponds to Execute method of Java's ExecutorService interface.
	Execute(task Task)

	TryExecute(task Task) bool

	Submit(groupID string, task Task)
	// Wait waits for all the tasks to complete.
	Wait(groupID string)

	Close()
}

type futureTask struct {
	task Task
	stop chan struct{}
}

func (f *futureTask) Run() {
	defer func() {
		if f.stop == nil {
			return
		}
		close(f.stop)
	}()
	defer func() {
		if err := recover(); err != nil {
			log.Error("run task failed", zap.String("err", fmt.Sprintf("%v", err)))
		}
	}()

	f.task.Run()
}

type executor struct {
	lock             sync.Mutex
	waitingTasks     map[string][]chan struct{}
	concurrencyLimit int

	workQueue chan futureTask
}

func (ex *executor) Execute(task Task) {
	ex.execute(task)
}

func (ex *executor) TryExecute(task Task) bool {
	return ex.tryExecute(task)
}

func (ex *executor) Submit(groupID string, task Task) {
	ex.submit(groupID, task)
}

func (ex *executor) Wait(groupID string) {
	ex.lock.Lock()
	defer ex.lock.Unlock()

	tasks, ok := ex.waitingTasks[groupID]
	if !ok {
		return
	}

	for _, t := range tasks {
		<-t
	}

	delete(ex.waitingTasks, groupID)
}

func (ex *executor) Close() {
	close(ex.workQueue)
}

func (ex *executor) execute(task Task) {
	ex.workQueue <- futureTask{
		task: task,
	}
}

func (ex *executor) tryExecute(task Task) bool {
	ex.lock.Lock()
	defer ex.lock.Unlock()

	var pushed bool
	ft := futureTask{task: task}
	select {
	case ex.workQueue <- ft:
		pushed = true
	default:
		pushed = false
	}

	return pushed
}

func (ex *executor) submit(groupID string, task Task) {
	ex.lock.Lock()
	defer ex.lock.Unlock()

	stopCh := make(chan struct{})

	waitingTasks, ok := ex.waitingTasks[groupID]
	if !ok {
		waitingTasks = make([]chan struct{}, 0)
	}

	waitingTasks = append(waitingTasks, stopCh)
	ex.waitingTasks[groupID] = waitingTasks
	ex.workQueue <- futureTask{
		task: task,
		stop: stopCh,
	}
}

func (ex *executor) runWorker() {
	for i := 0; i < ex.concurrencyLimit; i++ {
		go func() {
			for t := range ex.workQueue {
				if err := recover(); err != nil {
					log.Error("run task failed", zap.String("err", fmt.Sprintf("%v", err)))
				}

				t.Run()
			}
		}()
	}
}

// NewExecutor returns a new Executor.
func NewExecutor(concurrencyLimit int) Executor {
	ex := &executor{
		waitingTasks:     make(map[string][]chan struct{}),
		concurrencyLimit: concurrencyLimit,
		workQueue:        make(chan futureTask, concurrencyLimit),
	}

	ex.runWorker()
	return ex
}

func NewExecutorWithQueueSize(concurrencyLimit, queueSize int) Executor {
	ex := &executor{
		waitingTasks:     make(map[string][]chan struct{}),
		concurrencyLimit: concurrencyLimit,
		workQueue:        make(chan futureTask, queueSize),
	}

	ex.runWorker()
	return ex
}