From f2040886198b8f4a25ce1a9f65585097b299f9fa Mon Sep 17 00:00:00 2001
From: Jeff Carr <jcarr@wit.com>
Date: Tue, 22 Oct 2024 16:33:36 -0500
Subject: [PATCH] attempt at human readable message in JSON output

Signed-off-by: Jeff Carr <jcarr@wit.com>
---
 Makefile       |  3 ++-
 droplet.proto  |  7 ++++-
 helpers.go     | 71 ++++++++++++++++++++++++++++++++++++++------------
 sampleData.go  |  3 +--
 storageinfo.go | 53 +++++++++++++++++++++++++++++++++++++
 5 files changed, 117 insertions(+), 20 deletions(-)
 create mode 100644 storageinfo.go

diff --git a/Makefile b/Makefile
index 1d23b4c..e1a98e7 100644
--- a/Makefile
+++ b/Makefile
@@ -17,8 +17,9 @@ all:
 	make cluster.pb.go
 	make -C configfile
 
-vet: lint
+vet:
 	GO111MODULE=off go vet
+	make lint
 
 lint:
 	buf lint droplet.proto
diff --git a/droplet.proto b/droplet.proto
index eab564c..5af9db0 100644
--- a/droplet.proto
+++ b/droplet.proto
@@ -8,7 +8,7 @@ message Droplet {
 	string	name	= 2;
 	string	hostname = 3;
 	int64	cpus	= 4;
-	string	memory	= 5;
+	int64	memory	= 5;
 	int64	disk	= 6;
 	string	base_image = 7;
 
@@ -19,6 +19,7 @@ message Droplet {
 	string	default_state = 11;
 	string	preferred_hypervisor = 12;
 	google.protobuf.Any	testany	= 13;
+        StorageInfo testsi = 14;
 
 	message Network {
 		string	mac     = 1;
@@ -29,4 +30,8 @@ message Droplet {
 		string	filename        = 1;
 		int64	size    	= 2;
 	}
+
+        message StorageInfo {
+                int64 capacity = 1;  // Stores the storage capacity in bytes.
+        }
 }
diff --git a/helpers.go b/helpers.go
index 8f602ee..4ffd13c 100644
--- a/helpers.go
+++ b/helpers.go
@@ -8,6 +8,16 @@ import (
 	"google.golang.org/protobuf/encoding/prototext"
 )
 
+// can the json protobuf output use a string and have a type handler
+// to convert it back to int64?
+func SetGB(gb int) int64 {
+	return int64(gb * 1024 * 1024 * 1024)
+}
+
+func SetMB(mb int) int64 {
+	return int64(mb * 1024 * 1024)
+}
+
 func (x *Hypervisor) SetMemoryGB(gb int) {
 	x.Memory = int64(gb * 1024 * 1024 * 1024)
 }
@@ -36,22 +46,6 @@ func (c *Cluster) FormatTEXT() string {
 	return prototext.Format(c)
 }
 
-func (c *Cluster) AddHypervisor(hostname string, cpus int, mem int) *Hypervisor {
-	// Generate a new UUID
-	id := uuid.New()
-	h := &Hypervisor{
-		Uuid:     id.String(),
-		Hostname: hostname,
-		Cpus:     16,
-		Memory:   256,
-		Comment:  "this is a fake hypervisor",
-	}
-	h.Cpus = int64(cpus)
-	h.SetMemoryGB(mem * 32)
-	c.Hypervisors = append(c.Hypervisors, h)
-	return h
-}
-
 func (c *Cluster) FindDroplet(name string) *Droplet {
 	for _, d := range c.Droplets {
 		if d.Hostname == name {
@@ -69,3 +63,48 @@ func (c *Cluster) FindHypervisor(name string) *Hypervisor {
 	}
 	return nil
 }
+
+func (c *Cluster) AddHypervisor(hostname string, cpus int, mem int) *Hypervisor {
+	h := c.FindHypervisor(hostname)
+	if h != nil {
+		return h
+	}
+	// Generate a new UUID
+	id := uuid.New()
+	h = &Hypervisor{
+		Uuid:     id.String(),
+		Hostname: hostname,
+		Cpus:     int64(cpus),
+		Comment:  "this is a fake hypervisor",
+	}
+	if cpus < 0 {
+		h.Cpus = 1
+	}
+	h.SetMemoryGB(mem * 32)
+	c.Hypervisors = append(c.Hypervisors, h)
+	return h
+}
+
+func (c *Cluster) AddDroplet(hostname string, cpus int, mem int) *Droplet {
+	d := c.FindDroplet(hostname)
+	if d != nil {
+		return d
+	}
+
+	// Generate a new UUID
+	id := uuid.New()
+	d = &Droplet{
+		Uuid:     id.String(),
+		Hostname: hostname,
+		Cpus:     int64(cpus),
+	}
+
+	if cpus < 0 {
+		d.Cpus = 1
+	}
+	d.Memory = SetGB(mem * 32)
+	d.Testsi.Capacity = SetGB(mem * 32)
+	// d.Testsi = StorageInfo{Capacity: 64}
+	c.Droplets = append(c.Droplets, d)
+	return d
+}
diff --git a/sampleData.go b/sampleData.go
index e310dec..5a05377 100644
--- a/sampleData.go
+++ b/sampleData.go
@@ -51,8 +51,7 @@ func CreateSampleCluster(total int) *Cluster {
 		d.PreferredHypervisor = fmt.Sprintf("farm%d", i)
 		if d.PreferredHypervisor == "farm4" {
 			d.Cpus = 16
-			d.Memory = "256 MB"
-			// d.Memory = "258 MB"
+			d.Memory = SetGB(256)
 		}
 
 		c.Droplets = append(c.Droplets, d)
diff --git a/storageinfo.go b/storageinfo.go
new file mode 100644
index 0000000..9df42c2
--- /dev/null
+++ b/storageinfo.go
@@ -0,0 +1,53 @@
+package virtbuf
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+)
+
+type StorageInfo struct {
+	Capacity int64
+}
+
+// MarshalJSON custom marshals the StorageInfo struct to JSON
+func (s StorageInfo) MarshalJSON() ([]byte, error) {
+	capacityStr := fmt.Sprintf("%d GB", s.Capacity)
+	return json.Marshal(map[string]string{
+		"capacity": capacityStr,
+	})
+}
+
+// UnmarshalJSON custom unmarshals JSON into the StorageInfo struct
+func (s *StorageInfo) UnmarshalJSON(data []byte) error {
+	var raw map[string]string
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	if capacityStr, ok := raw["capacity"]; ok {
+		capacityStr = capacityStr[:len(capacityStr)-3] // Remove the " GB" suffix
+		capacity, err := strconv.ParseInt(capacityStr, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.Capacity = capacity
+	}
+
+	return nil
+}
+
+/*
+func main() {
+	info := StorageInfo{Capacity: 64}
+
+	// Marshaling to JSON
+	jsonData, _ := json.Marshal(info)
+	fmt.Println(string(jsonData)) // Output: {"capacity":"64 GB"}
+
+	// Unmarshaling back to Go struct
+	var newInfo StorageInfo
+	_ = json.Unmarshal(jsonData, &newInfo)
+	fmt.Println(newInfo.Capacity) // Output: 64
+}
+*/