package main

import (
	"errors"
	"sync"

	"github.com/barakmich/tinkerbell/ray_rpc"
	"go.uber.org/zap"
)

type WorkerPool interface {
	Schedule(*ray_rpc.Work) error
	Close() error
	Finish(*ray_rpc.WorkStatus) error
	Register(worker Worker) error
	Deregister(worker Worker) error
}

type SimpleRRWorkerPool struct {
	sync.Mutex
	workers []Worker
	store   ObjectStore
	offset  int
}

func NewRoundRobinWorkerPool(obj ObjectStore) *SimpleRRWorkerPool {
	return &SimpleRRWorkerPool{
		store: obj,
	}
}

func (wp *SimpleRRWorkerPool) Register(worker Worker) error {
	wp.Lock()
	defer wp.Unlock()
	wp.workers = append(wp.workers, worker)
	return nil
}

func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
	wp.Lock()
	defer wp.Unlock()
	if len(wp.workers) == 0 {
		return errors.New("No workers available, try again later")
	}
	origOffset := wp.offset

	done := false
	for !done {
		worker := wp.workers[wp.offset]
		if worker.Schedulable() {
			zap.S().Info("Sending work to worker", wp.offset)
			wp.workers[wp.offset].AssignWork(work)
			done = true
		}
		wp.offset++
		if wp.offset == len(wp.workers) {
			wp.offset = 0
		}
		if wp.offset == origOffset && !done {
			return errors.New("No workers schedulable")
		}
	}
	return nil
}

func (wp *SimpleRRWorkerPool) Finish(status *ray_rpc.WorkStatus) error {
	if status.Status != ray_rpc.COMPLETE {
		panic("todo: Only call Finish on successfully completed work")
	}
	id := deserializeObjectID(status.FinishedTicket.ReturnId)
	return wp.store.PutObject(&Object{id, status.CompleteData})
}

func (wp *SimpleRRWorkerPool) Close() error {
	wp.Lock()
	defer wp.Unlock()
	for _, w := range wp.workers {
		w.Close()
	}
	return nil
}

func (wp *SimpleRRWorkerPool) Deregister(worker Worker) error {
	wp.Lock()
	defer wp.Unlock()
	found := false
	for i, w := range wp.workers {
		if w == worker {
			wp.workers = append(wp.workers[:i], wp.workers[i+1:]...)
			if wp.offset == len(wp.workers) {
				wp.offset = 0
			}
			worker.Close()
			found = true
			zap.S().Info("Deregistering worker", i)
		}
	}
	if !found {
		panic("Trying to deregister a worker that was never created")
	}
	return nil
}