From 185e8cbb562a47b71bded30a4d3d7e2167d6dfcf Mon Sep 17 00:00:00 2001
From: John Mark Gabriel Caguicla <jmg.caguicla@guarandoo.me>
Date: Tue, 28 Jun 2022 05:06:50 +0800
Subject: [PATCH] add stock query commands

---
 api/client.go     | 43 ++++++++++++++++++++++
 commands/stock.go | 93 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+)
 create mode 100644 commands/stock.go

diff --git a/api/client.go b/api/client.go
index 02d7a05..848837b 100644
--- a/api/client.go
+++ b/api/client.go
@@ -58,6 +58,21 @@ type DeployServerResponse struct {
 	} `json:"server"`
 }
 
+type ListGpuStockResponse struct {
+	Response
+	Stock map[string]map[string]struct {
+		AvailableNow     int `json:"available_now"`
+		AvailableReserve int `json:"available_reserve"`
+	} `json:"stock"`
+}
+
+type ListCpuStockResponse struct {
+	Response
+	Stock map[string]map[string]struct {
+		AvailableNow string `json:"available_now"`
+	} `json:"stock"`
+}
+
 type Client struct {
 	BaseUrl  string
 	ApiKey   string
@@ -283,3 +298,31 @@ func (client *Client) RestartServer(server string) (*Response, error) {
 
 	return &res, nil
 }
+
+func (client *Client) ListGpuStock() (*ListGpuStockResponse, error) {
+	raw, err := client.get("stock/list", nil, false)
+	if err != nil {
+		return nil, err
+	}
+
+	var res ListGpuStockResponse
+	if err := json.Unmarshal(*raw, &res); err != nil {
+		return nil, err
+	}
+
+	return &res, nil
+}
+
+func (client *Client) ListCpuStock() (*ListCpuStockResponse, error) {
+	raw, err := client.get("stock/cpu/list", nil, false)
+	if err != nil {
+		return nil, err
+	}
+
+	var res ListCpuStockResponse
+	if err := json.Unmarshal(*raw, &res); err != nil {
+		return nil, err
+	}
+
+	return &res, nil
+}
diff --git a/commands/stock.go b/commands/stock.go
new file mode 100644
index 0000000..5d8d635
--- /dev/null
+++ b/commands/stock.go
@@ -0,0 +1,93 @@
+package commands
+
+import (
+	"errors"
+	"os"
+
+	"github.com/jedib0t/go-pretty/v6/table"
+	"github.com/spf13/cobra"
+)
+
+var (
+	stockCmd = &cobra.Command{
+		Use:   "stock",
+		Short: "Query stock",
+	}
+	listStockCmd = &cobra.Command{
+		Use:   "list",
+		Short: "List stock",
+		RunE:  listStock,
+	}
+)
+
+func init() {
+	stockCmd.AddCommand(listStockCmd)
+	listStockCmd.Flags().String("type", "gpu", "Instance type (gpu or cpu)")
+	listStockCmd.Flags().Bool("all", false, "Include out-of-stock instances")
+
+	rootCmd.AddCommand(stockCmd)
+}
+
+func listStock(cmd *cobra.Command, args []string) error {
+	instanceType, err := cmd.Flags().GetString("type")
+	if err != nil {
+		return err
+	}
+
+	all, err := cmd.Flags().GetBool("all")
+	if err != nil {
+		return err
+	}
+
+	t := table.NewWriter()
+	t.SetOutputMirror(os.Stdout)
+
+	switch instanceType {
+	case "gpu":
+		res, err := client.ListGpuStock()
+		if err != nil {
+			return nil
+		}
+
+		if !res.Success {
+			return errors.New(res.Error)
+		}
+
+		t.AppendHeader(table.Row{"GPU", "Region", "Available Now", "Available Reserve"})
+
+		for gpuModel, regionStock := range res.Stock {
+			for region, stock := range regionStock {
+				if (stock.AvailableNow > 0 || stock.AvailableReserve > 0) || all {
+					t.AppendRow(table.Row{gpuModel, region, stock.AvailableNow, stock.AvailableReserve})
+				}
+			}
+		}
+
+	case "cpu":
+		res, err := client.ListCpuStock()
+		if err != nil {
+			return nil
+		}
+
+		if !res.Success {
+			return errors.New(res.Error)
+		}
+
+		t.AppendHeader(table.Row{"CPU Model", "Region", "Available Now"})
+
+		for cpuModel, regionStock := range res.Stock {
+			for region, stock := range regionStock {
+				if stock.AvailableNow != "None" || all {
+					t.AppendRow(table.Row{cpuModel, region, stock.AvailableNow})
+				}
+			}
+		}
+
+	default:
+		return errors.New("unknown instance type")
+	}
+
+	t.Render()
+
+	return nil
+}