package middleware_http

import (
	"bytes"
	"encoding/json"
	"fmt"
	"gitee.com/os-lee/easy-paas/common"
	"gitee.com/os-lee/easy-paas/common/e"
	"gitee.com/os-lee/easy-paas/gateway/internal/dao"
	"gitee.com/os-lee/easy-paas/gateway/internal/load_balancer"
	"gitee.com/os-lee/easy-paas/gateway/internal/reverse_proxy"
	"github.com/gin-gonic/gin"
	"github.com/pkg/errors"
	"net/http"
)

// InterceptingResponseWriter 是一个自定义的 ResponseWriter，用于捕获响应数据。
type InterceptingResponseWriter struct {
	http.ResponseWriter               // 原始的 ResponseWriter
	body                *bytes.Buffer // 缓冲区，用于保存响应数据
}

// Write 方法重写，将数据同时写入原始的 ResponseWriter 和缓冲区。
func (irw *InterceptingResponseWriter) Write(b []byte) (int, error) {
	irw.body.Write(b)
	return irw.ResponseWriter.Write(b)
}

// Flush 方法，用于支持流式传输，如果底层的 ResponseWriter 支持的话。
func (irw *InterceptingResponseWriter) Flush() {
	if f, ok := irw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// ReverseProxy 创建并返回一个gin.HandlerFunc，
// 该中间件用于根据服务名动态发现后端服务，并通过负载均衡策略反向代理请求到相应的后端服务。
func ReverseProxy() gin.HandlerFunc {
	return func(c *gin.Context) {
		serverInterface, ok := c.Get("service")
		if !ok {
			common.ResponseError(c, e.ErrGateway, errors.New("从上下文中获取服务失败"))
			c.Abort()
			return
		}
		serviceDetail := serverInterface.(*dao.ServiceDetail)
		lb, err := load_balancer.Handler.Get(serviceDetail)
		if err != nil {
			common.ResponseError(c, e.ErrGateway,
				errors.New(fmt.Sprintf("获取负载均衡器错误, 服务名: %s Err: %v", serviceDetail.Info.ServiceName, err)))
			c.Abort()
			return
		}
		proxy := reverse_proxy.NewLoadBalanceReverseProxy(c, lb)

		// 创建一个InterceptingResponseWriter实例
		interceptingWriter := &InterceptingResponseWriter{
			ResponseWriter: c.Writer,
			body:           bytes.NewBuffer(nil),
		}

		// 执行代理请求
		proxy.ServeHTTP(interceptingWriter, c.Request)

		// 在代理请求之后，interceptingWriter.body 包含了响应数据
		responseBody := interceptingWriter.body.String()

		response, _ := json.Marshal(responseBody)
		c.Set("response", string(response))

		c.Abort()
		return
	}
}