package main

import (
	"context"

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

type Raylet struct {
	Objects ObjectStore
	Workers WorkerPool
}

func (r *Raylet) GetObject(_ context.Context, req *ray_rpc.GetRequest) (*ray_rpc.GetResponse, error) {
	zap.S().Debug("GetObject")
	data, err := GetObject(r.Objects, deserializeObjectID(req.Id))
	if err != nil {
		return nil, err
	}
	return &ray_rpc.GetResponse{
		Valid: true,
		Data:  data,
	}, nil
}

func (r *Raylet) PutObject(_ context.Context, req *ray_rpc.PutRequest) (*ray_rpc.PutResponse, error) {
	zap.S().Debug("PutObject")
	id := r.Objects.MakeID()
	err := r.Objects.PutObject(&Object{id, req.Data})
	if err != nil {
		return nil, err
	}
	return &ray_rpc.PutResponse{
		Id: serializeObjectID(id),
	}, nil
}

func (r *Raylet) WaitObject(_ context.Context, _ *ray_rpc.WaitRequest) (*ray_rpc.WaitResponse, error) {
	panic("not implemented") // TODO: Implement
}

func (r *Raylet) Schedule(_ context.Context, task *ray_rpc.ClientTask) (*ray_rpc.ClientTaskTicket, error) {
	zap.S().Debug("Schedule:", task.Type)
	id := r.Objects.MakeID()
	ticket := &ray_rpc.ClientTaskTicket{serializeObjectID(id)}
	work := &ray_rpc.Work{}
	work.Task = task
	work.Ticket = ticket
	err := r.Workers.Schedule(work)
	if err != nil {
		return nil, err
	}
	return ticket, nil
}

func (r *Raylet) Workstream(conn WorkstreamConnection) error {
	return r.Workers.Workstream(conn)
}

//func (r *Raylet) Workstream()

func NewMemRaylet() *Raylet {
	store := NewMemObjectStore()
	return &Raylet{
		Objects: store,
		Workers: NewRoundRobinWorkerPool(store),
	}
}