Moonshark/runner/env.go

package runner

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"Moonshark/utils/color"
	"Moonshark/utils/logger"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// EnvManager handles loading, storing, and saving environment variables
type EnvManager struct {
	envPath string         // Path to .env file
	vars    map[string]any // Environment variables in memory
	mu      sync.RWMutex   // Thread-safe access
}

// Global environment manager instance
var globalEnvManager *EnvManager

// InitEnv initializes the environment manager with the given data directory
func InitEnv(dataDir string) error {
	if dataDir == "" {
		return fmt.Errorf("data directory cannot be empty")
	}

	// Create data directory if it doesn't exist
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		return fmt.Errorf("failed to create data directory: %w", err)
	}

	envPath := filepath.Join(dataDir, ".env")

	globalEnvManager = &EnvManager{
		envPath: envPath,
		vars:    make(map[string]any),
	}

	// Load existing .env file if it exists
	if err := globalEnvManager.load(); err != nil {
		logger.Warning("Failed to load .env file: %v", err)
	}

	count := len(globalEnvManager.vars)
	if count > 0 {
		logger.Info("Environment loaded: %s vars from %s",
			color.Yellow(fmt.Sprintf("%d", count)),
			color.Yellow(envPath))
	} else {
		logger.Info("Environment initialized: %s", color.Yellow(envPath))
	}

	return nil
}

// GetGlobalEnvManager returns the global environment manager instance
func GetGlobalEnvManager() *EnvManager {
	return globalEnvManager
}

// load reads the .env file and populates the vars map
func (e *EnvManager) load() error {
	file, err := os.Open(e.envPath)
	if os.IsNotExist(err) {
		// File doesn't exist, start with empty env
		return nil
	}
	if err != nil {
		return fmt.Errorf("failed to open .env file: %w", err)
	}
	defer file.Close()

	e.mu.Lock()
	defer e.mu.Unlock()

	scanner := bufio.NewScanner(file)
	lineNum := 0

	for scanner.Scan() {
		lineNum++
		line := strings.TrimSpace(scanner.Text())

		// Skip empty lines and comments
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		// Parse key=value
		parts := strings.SplitN(line, "=", 2)
		if len(parts) != 2 {
			logger.Warning("Invalid .env line %d: %s", lineNum, line)
			continue
		}

		key := strings.TrimSpace(parts[0])
		value := strings.TrimSpace(parts[1])

		// Remove quotes if present
		if len(value) >= 2 {
			if (strings.HasPrefix(value, "\"") && strings.HasSuffix(value, "\"")) ||
				(strings.HasPrefix(value, "'") && strings.HasSuffix(value, "'")) {
				value = value[1 : len(value)-1]
			}
		}

		e.vars[key] = value
	}

	return scanner.Err()
}

// Save writes the current environment variables to the .env file
func (e *EnvManager) Save() error {
	if e == nil {
		return nil // No env manager initialized
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	file, err := os.Create(e.envPath)
	if err != nil {
		return fmt.Errorf("failed to create .env file: %w", err)
	}
	defer file.Close()

	// Sort keys for consistent output
	keys := make([]string, 0, len(e.vars))
	for key := range e.vars {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	// Write header comment
	fmt.Fprintln(file, "# Environment variables for Moonshark")
	fmt.Fprintln(file, "# Generated automatically - you can edit this file")
	fmt.Fprintln(file)

	// Write each variable
	for _, key := range keys {
		value := e.vars[key]

		// Convert value to string
		var strValue string
		switch v := value.(type) {
		case string:
			strValue = v
		case nil:
			continue // Skip nil values
		default:
			strValue = fmt.Sprintf("%v", v)
		}

		// Quote values that contain spaces or special characters
		if strings.ContainsAny(strValue, " \t\n\r\"'\\") {
			strValue = fmt.Sprintf("\"%s\"", strings.ReplaceAll(strValue, "\"", "\\\""))
		}

		fmt.Fprintf(file, "%s=%s\n", key, strValue)
	}

	logger.Debug("Environment saved: %d vars to %s", len(e.vars), e.envPath)
	return nil
}

// Get retrieves an environment variable
func (e *EnvManager) Get(key string) (any, bool) {
	if e == nil {
		return nil, false
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	value, exists := e.vars[key]
	return value, exists
}

// Set stores an environment variable
func (e *EnvManager) Set(key string, value any) {
	if e == nil {
		return
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	e.vars[key] = value
}

// GetAll returns a copy of all environment variables
func (e *EnvManager) GetAll() map[string]any {
	if e == nil {
		return make(map[string]any)
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	result := make(map[string]any, len(e.vars))
	for k, v := range e.vars {
		result[k] = v
	}
	return result
}

// CleanupEnv saves the environment and cleans up resources
func CleanupEnv() error {
	if globalEnvManager != nil {
		return globalEnvManager.Save()
	}
	return nil
}

// envGet Lua function to get an environment variable
func envGet(state *luajit.State) int {
	if !state.IsString(1) {
		state.PushNil()
		return 1
	}

	key := state.ToString(1)
	if value, exists := globalEnvManager.Get(key); exists {
		if err := state.PushValue(value); err != nil {
			state.PushNil()
		}
	} else {
		state.PushNil()
	}
	return 1
}

// envSet Lua function to set an environment variable
func envSet(state *luajit.State) int {
	if !state.IsString(1) || !state.IsString(2) {
		state.PushBoolean(false)
		return 1
	}

	key := state.ToString(1)
	value := state.ToString(2)

	globalEnvManager.Set(key, value)
	state.PushBoolean(true)
	return 1
}

// envGetAll Lua function to get all environment variables
func envGetAll(state *luajit.State) int {
	vars := globalEnvManager.GetAll()

	if err := state.PushTable(vars); err != nil {
		state.PushNil()
	}

	return 1
}

// RegisterEnvFunctions registers environment functions with the Lua state
func RegisterEnvFunctions(state *luajit.State) error {
	if err := state.RegisterGoFunction("__env_get", envGet); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__env_set", envSet); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__env_get_all", envGetAll); err != nil {
		return err
	}
	return nil
}