package media

import (
	"context"
	"errors"
	"time"

	"go.uber.org/zap"
)

// WorkerPool is a pool of workers that can consume and run a queue of tasks.
type WorkerPool struct {
	size   int
	ch     chan func()
	logger *zap.SugaredLogger
}

// NewWorkerPool returns a new WorkerPool containing the specified number of
// workers, and with the provided maximum queue size. Jobs added to the queue
// after it reaches this size limit will be rejected.
func NewWorkerPool(size int, maxQueueSize int, logger *zap.SugaredLogger) *WorkerPool {
	return &WorkerPool{
		size:   size,
		ch:     make(chan func(), maxQueueSize),
		logger: logger,
	}
}

// NewTestWorkerPool returns a new running WorkerPool with a single worker,
// and noop logger, suitable for test environments.
func NewTestWorkerPool() *WorkerPool {
	p := NewWorkerPool(1, 256, zap.NewNop().Sugar())
	p.Run()
	return p
}

// Run launches the workers, and returns immediately.
func (p *WorkerPool) Run() {
	for i := 0; i < p.size; i++ {
		go func() {
			for task := range p.ch {
				task()
			}
		}()
	}
}

// WaitForTask blocks while the provided task is executed by a worker,
// returning the error returned by the task.
func (p *WorkerPool) WaitForTask(ctx context.Context, taskFunc func() error) error {
	done := make(chan error)
	queuedAt := time.Now()
	fn := func() {
		startedAt := time.Now()
		result := taskFunc()

		durTotal := time.Since(queuedAt)
		durTask := time.Since(startedAt)
		durQueue := startedAt.Sub(queuedAt)
		p.logger.With("task", durTask, "queue", durQueue, "total", durTotal).Infof("Completed task")

		done <- result
	}

	select {
	case p.ch <- fn:
		return <-done
	case <-ctx.Done():
		return ctx.Err()
	default:
		return errors.New("worker queue full")
	}
}