main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	compute "cloud.google.com/go/compute/apiv1"
	"cloud.google.com/go/compute/apiv1/computepb"
	"go.uber.org/zap"
	"google.golang.org/api/option"
)

type Config struct {
	Port           string
	Zone           string
	Project        string
	VMNames        []string
	CredentialsKey string
	Logger         *zap.SugaredLogger
}

// ActiveJobWebhook represents the payload sent to customer webhooks
type ActiveJobWebhook struct {
	Event string `json:"event"`
	Owner string `json:"owner"`

	Timestamp time.Time `json:"timestamp"`
}

type VMBooter struct {
	config *Config
	client *compute.InstancesClient
	logger *zap.SugaredLogger
}

func main() {
	// Initialize logger
	logger, _ := zap.NewProduction()
	defer logger.Sync()
	sugar := logger.Sugar()

	// Load configuration
	config := loadConfig(sugar)

	// Initialize VM manager
	vmManager, err := NewVMBooter(config)
	if err != nil {
		sugar.Fatalw("Failed to initialize VM manager", "error", err)
	}
	defer vmManager.Close()

	// Setup HTTP server
	http.HandleFunc("/webhook", vmManager.handleWebhook)
	http.HandleFunc("/health", handleHealth)

	sugar.Infow("Starting server",
		"port", config.Port,
		"zone", config.Zone,
		"project", config.Project,
		"vms", config.VMNames,
	)

	if err := http.ListenAndServe(":"+config.Port, nil); err != nil {
		sugar.Fatalw("Server failed to start", "error", err)
	}
}

func loadConfig(logger *zap.SugaredLogger) *Config {
	port := getEnvOrDefault("PORT", "8080")
	zone := getEnvOrDefault("ZONE", "")
	project := getEnvOrDefault("PROJECT", "")
	vmNamesStr := getEnvOrDefault("VM_NAMES", "")
	credentialsKey := getEnvOrDefault("GOOGLE_APPLICATION_CREDENTIALS_JSON", "")

	if zone == "" {
		logger.Fatal("ZONE environment variable is required")
	}
	if project == "" {
		logger.Fatal("PROJECT environment variable is required")
	}
	if vmNamesStr == "" {
		logger.Fatal("VM_NAMES environment variable is required")
	}

	vmNames := strings.Split(vmNamesStr, ",")
	for i, name := range vmNames {
		vmNames[i] = strings.TrimSpace(name)
	}

	return &Config{
		Port:           port,
		Zone:           zone,
		Project:        project,
		VMNames:        vmNames,
		CredentialsKey: credentialsKey,
		Logger:         logger,
	}
}

func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

func NewVMBooter(config *Config) (*VMBooter, error) {
	ctx := context.Background()

	var client *compute.InstancesClient
	var err error

	// If running outside GCP, use static credentials
	if config.CredentialsKey != "" {
		config.Logger.Info("Using static credentials from GOOGLE_APPLICATION_CREDENTIALS_JSON")
		client, err = compute.NewInstancesRESTClient(ctx, option.WithCredentialsJSON([]byte(config.CredentialsKey)))
	} else {
		// Use inherited VM permissions (default credentials from metadata service)
		config.Logger.Info("Using default credentials (VM service account or ambient environment)")
		client, err = compute.NewInstancesRESTClient(ctx)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to create compute client: %w", err)
	}

	config.Logger.Info("Successfully initialized GCP Compute client")
	return &VMBooter{
		config: config,
		client: client,
		logger: config.Logger,
	}, nil
}

func (vm *VMBooter) Close() {
	if vm.client != nil {
		vm.client.Close()
	}
}

func (vm *VMBooter) handleWebhook(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var webhook ActiveJobWebhook
	if err := json.NewDecoder(r.Body).Decode(&webhook); err != nil {
		vm.logger.Errorw("Failed to decode webhook request", "error", err)
		http.Error(w, "Invalid JSON", http.StatusBadRequest)
		return
	}

	vm.logger.Infow("Received webhook", "event", webhook.Event, "owner", webhook.Owner, "timestamp", webhook.Timestamp)

	// Start all configured VMs
	for _, vmName := range vm.config.VMNames {
		if err := vm.startVM(vmName); err != nil {
			vm.logger.Errorw("Failed to start VM", "vm", vmName, "error", err)
			http.Error(w, fmt.Sprintf("Failed to start VM %s: %v", vmName, err), http.StatusInternalServerError)
			return
		}
	}

	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(map[string]string{
		"status":  "success",
		"message": "Processed event",
	})
}

func (vm *VMBooter) startVM(vmName string) error {
	ctx := context.Background()

	// First, get the current status of the VM
	getInstance := &computepb.GetInstanceRequest{
		Project:  vm.config.Project,
		Zone:     vm.config.Zone,
		Instance: vmName,
	}

	instance, err := vm.client.Get(ctx, getInstance)
	if err != nil {
		return fmt.Errorf("failed to get VM status: %w", err)
	}

	currentStatus := instance.GetStatus()

	// If VM is already running, just log and return
	if currentStatus == "RUNNING" {
		vm.logger.Infow("VM is already running", "vm", vmName)
		return nil
	}

	// If VM is stopped, start it
	if currentStatus == "TERMINATED" || currentStatus == "STOPPED" {
		startReq := &computepb.StartInstanceRequest{
			Project:  vm.config.Project,
			Zone:     vm.config.Zone,
			Instance: vmName,
		}

		op, err := vm.client.Start(ctx, startReq)
		if err != nil {
			return fmt.Errorf("failed to start VM: %w", err)
		}

		vm.logger.Infow("VM start initiated",
			"vm", vmName,
			"status", "starting",
		)

		st := time.Now()
		go func() {
			ctx := context.Background()
			op.Wait(ctx)
			vm.logger.Infow("VM was started successfully", "vm", vmName, "duration", time.Since(st))
		}()

		return nil
	}

	// VM is in some other state (STAGING, STOPPING, etc.)
	vm.logger.Infow("VM is in transitional state", "vm", vmName, "status", currentStatus)
	return nil
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(map[string]string{
		"status": "healthy",
	})
}