package response

import (
	"bytes"

	"braces.dev/errtrace"
	"github.com/gin-gonic/gin"
)

type WriterWrapper struct {
	gin.ResponseWriter

	Body *bytes.Buffer
}

func (w *WriterWrapper) Write(b []byte) (int, error) {
	// double write
	return errtrace.Wrap2(w.Body.Write(b))
}

func (w *WriterWrapper) WriteString(s string) (int, error) {
	// double write
	return errtrace.Wrap2(w.Body.WriteString(s))
}

func WriterWrapperHandler(ctx *gin.Context) {
	w := &WriterWrapper{
		Body:           &bytes.Buffer{},
		ResponseWriter: ctx.Writer,
	}
	ctx.Writer = w

	ctx.Next()

	// Write the response to gin response writer
	// if the gin response writer is empty
	_, _ = w.ResponseWriter.Write(w.Body.Bytes())
}