package main

import (
	"bytes"
	"context"
	"io/ioutil"
	"log"
	"net/http"
)

// Define a key type to avoid context key collisions.
type contextKey string

const bufferedBodyKey = contextKey("bufferedBody")

// bufferBodyMiddleware reads the request body and replaces it with a new reader,
// allowing it to be read multiple times. The original body is stored in the request context.
func bufferBodyMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only buffer if there's a body to read.
		if r.Body == nil || r.ContentLength == 0 {
			next.ServeHTTP(w, r)
			return
		}

		bodyBytes, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Printf("Error reading body in middleware: %v\n", err)
			return
		}
		defer r.Body.Close()

		// Store the buffered body in the context for downstream handlers.
		ctx := context.WithValue(r.Context(), bufferedBodyKey, bodyBytes)

		// Replace the original body with a new reader on the buffered bytes.
		// This allows subsequent handlers to read the body again.
		r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))

		// Call the next handler in the chain with the modified request.
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

/*
// okHandler is the final handler. It can now safely access the body from the context,
// knowing that other middleware might have also read it.
func okHandler(w http.ResponseWriter, r *http.Request) {
	// For demonstration, we can try reading the body directly here too.
	// The middleware ensures this is a fresh stream of the buffered data.
	bodyFromStream, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Printf("Error reading body in handler: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	log.Printf("Handler read %d bytes from the request body stream.", len(bodyFromStream))

	// We can also retrieve the body from the context if needed.
	bodyFromContext, ok := r.Context().Value(bufferedBodyKey).([]byte)
	if !ok {
		log.Println("Could not retrieve buffered body from context.")
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	log.Printf("Handler retrieved %d bytes from context.", len(bodyFromContext))

	// Prove they are the same.
	if !bytes.Equal(bodyFromStream, bodyFromContext) {
		log.Println("FATAL: Body from stream and context do not match!")
	}

	fmt.Fprintf(w, "Successfully read body of %d bytes.\n", len(bodyFromContext))
}
*/