// https://developer.fastly.com/solutions/examples/follow-redirects-at-the-edge
package main

import (
	"context"
	"fmt"
	"io"
	"regexp"

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

// Backend is the name of our primary service backend.
const Backend = "origin_0"

// MaxRedirectCount is the max number of redirects allowed.
const MaxRedirectCount = 2

// LocationHostPath captures a possible host and path from a Location header.
var LocationHostPath = regexp.MustCompile(`^(?:https?://([^/]+))?(/.*)?$`)

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		// Send the request to the named backend.
		resp, err := r.Send(ctx, Backend)
		if err != nil {
			w.WriteHeader(fsthttp.StatusBadGateway)
			fmt.Fprintln(w, err)
			return
		}

    // A request can't call .Send() twice.
		// The clone will be used as a clean base to clone from in the following loop.
		baseReq := r.Clone()

		var redirectCount int
		for isRedirect(resp.StatusCode) && redirectCount <= MaxRedirectCount {
			redirectCount++

			// A request can't call .Send() twice.
			// We clone the base request so we can follow any redirects.
			rc := baseReq.Clone()

			// Reset the URL path.
      // NOTE: This solution works for same-origin redirects only.
			rc.URL.Path = parsePath(resp.Header.Get("Location"))
			fmt.Printf("Redirect URL (count: %d): %+v\n", redirectCount, rc.URL)

			// Send the request to the named backend.
			resp, err = rc.Send(ctx, Backend)
			if err != nil {
				w.WriteHeader(fsthttp.StatusBadGateway)
				fmt.Fprintln(w, err)
				return
			}
		}

		// Write response to client.
		flush(resp, w)
	})
}

func isRedirect(statusCode int) bool {
	return statusCode >= 300 && statusCode < 400
}

func parsePath(location string) (path string) {
	ss := LocationHostPath.FindStringSubmatch(location)

	// NOTE: ss[0] is the full matched expression.
	if len(ss) >= 3 {
		path = ss[2]
	}

	// Set fallback if no path is matched.
	if path == "" {
		path = "/"
	}
	return path
}

func flush(resp *fsthttp.Response, w fsthttp.ResponseWriter) {
	w.Header().Reset(resp.Header)
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}