package main

import (
	"errors"

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

type WorkstreamConnection = ray_rpc.RayletWorkerConnection_WorkstreamServer

type Worker interface {
	AssignWork(work *ray_rpc.Work) error
	Run() error
	Close() error
	Schedulable() bool
}

type SimpleWorker struct {
	workChan   chan *ray_rpc.Work
	clientConn WorkstreamConnection
	pool       WorkerPool
	max        int
	curr       int
}

func (s *SimpleWorker) Schedulable() bool {
	return s.curr < s.max
}

func (s *SimpleWorker) AssignWork(work *ray_rpc.Work) error {
	s.workChan <- work
	return nil
}

func (s *SimpleWorker) Close() error {
	close(s.workChan)
	return nil
}

func (w *SimpleWorker) Run() error {
	sentinel, err := w.clientConn.Recv()
	if err != nil {
		return err
	}
	if sentinel.Status != ray_rpc.READY {
		return errors.New("Sent wrong sentinel? Closing...")
	}
	zap.S().Info("New worker:", sentinel.ErrorMsg)
	go func() {
		for work := range w.workChan {
			zap.S().Debug("Sending work")
			w.curr++
			err = w.clientConn.Send(work)
			if err != nil {
				zap.S().Error("Error sending:", err)
				return
			}
		}
	}()
	for {
		result, err := w.clientConn.Recv()
		if err != nil {
			zap.S().Error("Error on channel:", err)
			return err
		}
		w.curr--
		err = w.pool.Finish(result)
		if err != nil {
			zap.S().Error("Error finishing:", err)
			return err
		}
	}
	return nil
}