package device

import (
	"context"
	"time"

	pb "git.pengzhan.dev/ddp-simulator/proto"
	log "github.com/sirupsen/logrus"
)

type Device struct {
	pb.UnimplementedDeviceServer
	ID         string
	Role       string
	IP         string
	Port       int
	ParentID   *string
	ParentIP   string
	ParentPort int
	Capacity   *float32

	Children map[string]*Device
}

func NewDevice(id, role, ip string, port int) *Device {
	return &Device{
		ID:       id,
		Role:     role,
		IP:       ip,
		Port:     port,
		Children: make(map[string]*Device),
	}
}

func (d *Device) UpdateCapacity(capacity *float32) {
	d.Capacity = capacity
}

func (d *Device) UpdateParent(parentID *string, parentIP string, parentPort int) {
	d.ParentID = parentID
	d.ParentIP = parentIP
	d.ParentPort = parentPort
}

func (d *Device) Register(ctx context.Context, in *pb.RegisterRequest) (*pb.RegisterResponse, error) {
	child := NewDevice(in.Id, in.Role, in.Ip, int(in.Port))
	if in.Capacity != nil {
		child.UpdateCapacity(in.Capacity)
	}
	d.Children[in.Id] = child
	log.WithField("device", d.ID).Infof("device %s claims on device %s as its child", child.ID, d.ID)

	// Update with parent
	if d.ParentID != nil {
		if err := d.RegisterClientRetry(ctx, 30*time.Second); err != nil {
			log.WithField("device", d.ID).Fatalf("failed to register: %v", err)
		}
	}
	return &pb.RegisterResponse{}, nil
}