From 19f632535760dc03def880ee1cceebd6cba4195c Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Tue, 16 Jul 2024 22:47:24 +0000
Subject: [PATCH] cluster: rework `env` output to qualify zones using the
 region

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/cluster/env.go | 84 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 65 insertions(+), 19 deletions(-)

diff --git a/pkg/cluster/env.go b/pkg/cluster/env.go
index 42e98a2..a00964a 100644
--- a/pkg/cluster/env.go
+++ b/pkg/cluster/env.go
@@ -7,6 +7,8 @@ import (
 	"sort"
 	"strings"
 
+	"darvaza.org/core"
+
 	"git.jpi.io/amery/jpictl/pkg/rings"
 )
 
@@ -36,15 +38,19 @@ func (m *Cluster) Env(export bool) (*Env, error) {
 	return env, nil
 }
 
-// Zones returns the list of Zone IDs
-func (m *Env) Zones() []rings.ZoneID {
+// Zones returns the list of Zone IDs of a region,
+// or from all if none is specified.
+func (m *Env) Zones(r *Region) []rings.ZoneID {
 	var zones []rings.ZoneID
 
-	m.ForEachZone(func(z *Zone) bool {
+	iter := core.IIf[ZoneIterator](r != nil, r, m)
+
+	iter.ForEachZone(func(z *Zone) bool {
 		zones = append(zones, z.ID)
 		return false
 	})
 
+	core.SliceSort(zones, cmpOrdered[rings.ZoneID])
 	return zones
 }
 
@@ -64,6 +70,19 @@ func (m *Env) Regions() []string {
 	return regions
 }
 
+func (m *Env) RegionsID() (regions []rings.RegionID) {
+	m.ForEachRegion(func(r *Region) bool {
+		if r.IsPrimary() {
+			regions = append(regions, r.ID)
+		}
+
+		return false
+	})
+
+	core.SliceSort(regions, cmpOrdered[rings.RegionID])
+	return regions
+}
+
 // WriteTo generates environment variables for shell scripts
 func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
@@ -73,40 +92,55 @@ func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	}
 
 	m.writeEnvVar(&buf, genEnvStrings(m.Regions()), "REGIONS")
-	m.writeEnvVar(&buf, genEnvInts(m.Zones()), "ZONES")
+	m.writeEnvVar(&buf, genEnvInts(m.RegionsID()), "REGIONS_ID")
 
-	m.ForEachZone(func(z *Zone) bool {
-		m.writeEnvZone(&buf, z)
+	m.ForEachRegion(func(r *Region) bool {
+		if r.IsPrimary() {
+			m.writeEnvRegion(&buf, r)
+		}
 		return false
 	})
 
 	return buf.WriteTo(w)
 }
 
-func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
-	zoneID := z.ID
+func (m *Env) writeEnvRegion(w io.Writer, r *Region) {
+	regionID := r.ID
+
+	// REGION{regionID}_NAME
+	m.writeEnvVar(w, r.Name, "REGION%v_%s", regionID, "NAME")
+
+	// REGION{regionID}_ZONES
+	m.writeEnvVar(w, genEnvInts(m.Zones(r)), "REGION%v_%s", regionID, "ZONES")
+
+	r.ForEachZone(func(z *Zone) bool {
+		m.writeEnvZone(w, r, z)
+		return false
+	})
+}
+
+func (m *Env) writeEnvZone(w io.Writer, r *Region, z *Zone) {
+	zonePrefix := fmt.Sprintf("REGION%v_ZONE%v", r.ID, z.ID)
+	monPrefix := zonePrefix + "_MON"
 
 	// ZONE{zoneID}
-	m.writeEnvVar(w, genEnvZoneNodes(z), "ZONE%v", zoneID)
+	m.writeEnvVar(w, genEnvZoneNodes(z), zonePrefix)
 
 	// ZONE{zoneID}_NAME
-	m.writeEnvVar(w, z.Name, "ZONE%v_%s", zoneID, "NAME")
+	m.writeEnvVar(w, z.Name, zonePrefix+"_NAME")
 
 	// ZONE{zoneID}_GW
 	gateways, _ := z.GatewayIDs()
-	m.writeEnvVar(w, genEnvInts(gateways), "ZONE%v_%s", zoneID, "GW")
-
-	// ZONE{zoneID}_REGION
-	m.writeEnvVar(w, genEnvZoneRegion(z), "ZONE%v_%s", zoneID, "REGION")
+	m.writeEnvVar(w, genEnvInts(gateways), zonePrefix+"_GW")
 
 	// Ceph
 	monitors := z.GetCephMonitors()
 	// MON{zoneID}_NAME
-	m.writeEnvVar(w, genEnvZoneCephMonNames(monitors), "MON%v_%s", zoneID, "NAME")
+	m.writeEnvVar(w, genEnvZoneCephMonNames(monitors), monPrefix)
 	// MON{zoneID}_IP
-	m.writeEnvVar(w, genEnvZoneCephMonIPs(monitors), "MON%v_%s", zoneID, "IP")
+	m.writeEnvVar(w, genEnvZoneCephMonIPs(monitors), monPrefix+"_IP")
 	// MON{zoneID}_ID
-	m.writeEnvVar(w, genEnvZoneCephMonIDs(monitors), "MON%v_%s", zoneID, "ID")
+	m.writeEnvVar(w, genEnvZoneCephMonIDs(monitors), monPrefix+"_ID")
 }
 
 func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
@@ -131,14 +165,15 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	}
 }
 
-func genEnvInts[T ~int | ~uint](values []T) string {
+func genEnvInts[T core.Signed](values []T) string {
 	var buf bytes.Buffer
 
 	for _, v := range values {
 		if buf.Len() > 0 {
 			_, _ = buf.WriteRune(' ')
 		}
-		_, _ = buf.WriteString(fmt.Sprintf("%v", v))
+
+		_, _ = buf.WriteString(fmt.Sprintf("%v", int64(v)))
 	}
 
 	return buf.String()
@@ -209,3 +244,14 @@ func genEnvZoneCephMonIDs(m Machines) string {
 	})
 	return buf.String()
 }
+
+func cmpOrdered[T core.Ordered](a, b T) int {
+	switch {
+	case a == b:
+		return 0
+	case a < b:
+		return -1
+	default:
+		return 1
+	}
+}