package main

import (
	"context"
	"fmt"
	"io"
	"math/rand"

	"github.com/fastly/compute-sdk-go/fsthttp"
)

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		backends := []string{"origin_0", "origin_1"}

		// 1. Shuffle the backends slice randomly in-place
		rand.Shuffle(len(backends), func(i, j int) {
			backends[i], backends[j] = backends[j], backends[i]
		})

		var chosenBackend string

		// 2. Loop through the randomized backends and find the first healthy one
		for _, name := range backends {
			backend, err := fsthttp.BackendFromName(name)
			if err != nil {
				fmt.Printf("Backend not found: %s\n", name)
				continue
			}
			health, err := backend.Health()
			if err == nil && health == fsthttp.BackendHealthHealthy {
				fmt.Printf("Found healthy backend: %s\n", name)
				chosenBackend = name
				break
			}

			fmt.Printf("Backend %s is not explicitly healthy, trying next...\n", name)
		}

		// 3. Handle the fallback if absolutely no backends are healthy
		if chosenBackend == "" {
			fmt.Println("All backends are unhealthy or unknown!")
			w.WriteHeader(fsthttp.StatusServiceUnavailable)
			fmt.Fprintln(w, "Service Unavailable: No healthy backends available.")
			return
		}

		// 4. Send the request to the selected healthy backend
		fmt.Println("Sending request...")
		resp, err := r.Send(ctx, chosenBackend)
		if err != nil {
			w.WriteHeader(fsthttp.StatusBadGateway)
			fmt.Fprintln(w, err.Error())
			return
		}

		// 5. Stream response back to client
		w.Header().Reset(resp.Header)
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	})
}