package main

import (
	"context"
	"fmt"
	"io"
  "net/url"

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

// BackendName is the name of our service backend.
const BackendName = "http_me"

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		resp, err := r.Send(ctx, BackendName)
		if err != nil {
			w.WriteHeader(fsthttp.StatusBadGateway)
			fmt.Fprintln(w, err.Error())
			return
		}

    statusCode := resp.StatusCode;
    fmt.Println("Status code:", statusCode)
    
    switch statusCode {
      // Only redirect if status code is 301, 302 or 308.
      case fsthttp.StatusMovedPermanently, fsthttp.StatusFound, fsthttp.StatusPermanentRedirect:
        // Build a dynamic backend using the Location response header.
        location := resp.Header.Get("location")
        fmt.Println("Redirected to", location)
        // Parse the location to get the host name and the http scheme.
        u, err := url.Parse(location)
        if err != nil {
				fsthttp.Error(w, err.Error(), fsthttp.StatusInternalServerError)
			  }
        // Define backend options
        opts := fsthttp.NewBackendOptions()
        opts.HostOverride(u.Host)
        if u.Scheme == "https" {
          opts.UseSSL(true)
        } else {
          opts.UseSSL(false)
        }
        // Register dynamic backend
        fsthttp.RegisterDynamicBackend(u.Host, u.Host, opts)
        // Clone request 
        newReq := r.Clone()
        // Overwrite URL with redirect URL
        newReq.URL = u
        
        resp, err = newReq.Send(ctx, u.Host)
        if err != nil {
          w.WriteHeader(fsthttp.StatusBadGateway)
          fmt.Fprintln(w, err)
          return
        }

        w.Header().Reset(resp.Header)
        w.WriteHeader(resp.StatusCode)
        io.Copy(w, resp.Body)

      default: 
        // Do nothing. Return response as usual.
    }
		w.Header().Reset(resp.Header)
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	})
}