// https://developer.fastly.com/solutions/examples/add-or-remove-cookies
package main

import (
	"context"
	"fmt"
	"io"

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

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

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		// Read a specific cookie
		c, err := r.Cookie("myCookie")
		if err != nil {
			w.WriteHeader(fsthttp.StatusBadGateway)
			fmt.Fprintln(w, err)
			return
		}

		// Remove cookie from the request.
		FilterClientCookies(r, "myCookie")

		// 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
		}

		// Set a cookie in the response
		fsthttp.SetCookie(resp.Header, &fsthttp.Cookie{
			Name:   "myCookie",
			Value:  "foo",
			Path:   "/",
			MaxAge: 60,
		})

		// You can set multiple cookies in one response
		resp.Header.Add("Set-Cookie", "mySecondCookie=bar; httpOnly")

		// It is usually a good idea to prevent downstream caching of
		// responses that set cookies
		resp.Header.Set("Cache-Control", "no-store, private")

		fmt.Println(resp.Header.Keys()) // set-cookie is present

		// Write response to client.
		w.Header().Reset(resp.Header)
		w.WriteHeader(resp.StatusCode)

		// NOTE: We can't write to `w` after the first read of the `myCookie` Cookie.
		// Otherwise it would cause all other header modifications to be ignored.
		//
		// Refer to the documentation:
		// https://pkg.go.dev/net/http@go1.18.3#ResponseWriter
		if c.Value != "" {
			fmt.Fprintf(w, "The value of myCookie is %q\n", c.Value)
		}

		io.Copy(w, resp.Body)
	})
}

// FilterClientCookies removes a list of cookies from the "Cookie" request header by name.
func FilterClientCookies(r *fsthttp.Request, names ...string) {
	filter := make(map[string]struct{}, len(names))
	for _, name := range names {
		filter[name] = struct{}{}
	}

	cookies := r.Cookies()
	r.Header.Del("Cookie")

	for _, c := range cookies {
		if _, ok := filter[c.Name]; ok {
			continue
		}
		r.Header.Add("Cookie", fmt.Sprintf("%s=%s", c.Name, c.Value))
	}
}