package toolgin

import (
	"bytes"
	"github.com/gin-gonic/gin"
)

// Gin 框架在中间件中获取 response body 的方法
// https://www.toutiao.com/article/7182178685490004484/?wid=1684317766913

// ResponseWriter 自定义一个结构体，实现 gin.ResponseWriter interface
// 用于获取响应body,重写响应body。获取响应code,重写响应code
type ResponseWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
	code int
}

// 重写 Write([]byte) (int, error) 方法
func (w *ResponseWriter) Write(b []byte) (int, error) {
	return w.body.Write(b) //向一个bytes.buffer中写一份数据来为获取body使用
}

// WriteHeader 重写 WriteHeader(statusCode int) 方法
func (w *ResponseWriter) WriteHeader(statusCode int) {
	w.code = statusCode //向一个code中写一份数据来为获取code使用
}

// GetBody 获取body
func (w *ResponseWriter) GetBody() *bytes.Buffer {
	return w.body
}

// GetCode 获取code
func (w *ResponseWriter) GetCode() int {
	return w.code
}

// CleanBody 清空body,可用与重写body
func (w *ResponseWriter) CleanBody() {
	w.body = bytes.NewBuffer([]byte{})
}

// NewResponseWriter 包装 gin.ResponseWriter 接口
func NewResponseWriter(c *gin.Context) *ResponseWriter {
	writer := ResponseWriter{
		ResponseWriter: c.Writer,
		body:           bytes.NewBuffer([]byte{}),
		code:           0,
	}
	c.Writer = &writer
	return &writer
}