From a32b0347c03d4256e667343deab314a4e60dc311 Mon Sep 17 00:00:00 2001
From: Darren Shepherd <darren@acorn.io>
Date: Fri, 2 Aug 2024 14:01:09 -0700
Subject: [PATCH] chore: add getenv

---
 gptscript.go      | 28 ++++++++++++++++++++++++++
 gptscript_test.go | 51 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/gptscript.go b/gptscript.go
index 24bc35b..d16199c 100644
--- a/gptscript.go
+++ b/gptscript.go
@@ -2,7 +2,10 @@ package gptscript
 
 import (
 	"bufio"
+	"bytes"
+	"compress/gzip"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -290,3 +293,28 @@ func determineProperCommand(dir, bin string) string {
 	slog.Debug("Using gptscript binary: " + bin)
 	return bin
 }
+
+func GetEnv(key, def string) string {
+	v := os.Getenv(key)
+	if v == "" {
+		return def
+	}
+
+	if strings.HasPrefix(v, `{"_gz":"`) && strings.HasSuffix(v, `"}`) {
+		data, err := base64.StdEncoding.DecodeString(v[8 : len(v)-2])
+		if err != nil {
+			return v
+		}
+		gz, err := gzip.NewReader(bytes.NewBuffer(data))
+		if err != nil {
+			return v
+		}
+		strBytes, err := io.ReadAll(gz)
+		if err != nil {
+			return v
+		}
+		return string(strBytes)
+	}
+
+	return v
+}
diff --git a/gptscript_test.go b/gptscript_test.go
index 05449da..0cf8c19 100644
--- a/gptscript_test.go
+++ b/gptscript_test.go
@@ -1018,3 +1018,54 @@ func TestGetCommand(t *testing.T) {
 		})
 	}
 }
+
+func TestGetEnv(t *testing.T) {
+	// Cleaning up
+	defer func(currentEnvValue string) {
+		os.Setenv("testKey", currentEnvValue)
+	}(os.Getenv("testKey"))
+
+	// Tests
+	testCases := []struct {
+		name           string
+		key            string
+		def            string
+		envValue       string
+		expectedResult string
+	}{
+		{
+			name:           "NoValueUseDefault",
+			key:            "testKey",
+			def:            "defaultValue",
+			envValue:       "",
+			expectedResult: "defaultValue",
+		},
+		{
+			name:           "ValueExistsNoCompress",
+			key:            "testKey",
+			def:            "defaultValue",
+			envValue:       "testValue",
+			expectedResult: "testValue",
+		},
+		{
+			name:     "ValueExistsCompressed",
+			key:      "testKey",
+			def:      "defaultValue",
+			envValue: `{"_gz":"H4sIAEosrGYC/ytJLS5RKEvMKU0FACtB3ewKAAAA"}`,
+
+			expectedResult: "test value",
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			os.Setenv(test.key, test.envValue)
+
+			result := GetEnv(test.key, test.def)
+
+			if result != test.expectedResult {
+				t.Errorf("expected: %s, got: %s", test.expectedResult, result)
+			}
+		})
+	}
+}