package templatecache

import (
	"bytes"
	"github.com/gin-gonic/gin"
	"github.com/patrickmn/go-cache"
	"net/http"
	"time"
)

var (
	Cache *cache.Cache
)

func init() {
	Cache = cache.New(5*time.Minute, 5*time.Minute)
}

// 自定义 ResponseWriter 用于捕获 HTML 输出
type templateResponseWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w *templateResponseWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

func WithCache(duration time.Duration) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		// 只缓存 GET 请求（HEAD 由 _web 层处理，_server 不会收到 HEAD）
		if ctx.Request.Method != "GET" {
			ctx.Next()
			return
		}

		// 生成缓存键（包含路径和查询参数）
		cacheKey := ctx.Request.URL.String()
		var cacheJson []byte
		if cacheData, exists := Cache.Get(cacheKey); exists {
			if item, ok := cacheData.([]byte); ok {
				cacheJson = item
			}
		}
		// 检查缓存
		if cacheJson != nil {
			ctx.Data(http.StatusOK, "application/json; charset=utf-8", cacheJson)
			ctx.Abort()
			return
		}

		writer := &templateResponseWriter{
			body:           bytes.NewBuffer(nil),
			ResponseWriter: ctx.Writer,
		}
		ctx.Writer = writer

		ctx.Next()

		// 只有成功的响应才缓存
		status := ctx.Writer.Status()
		if status == 200 {
			cacheJsonData := writer.body.Bytes()
			cacheDuration := duration
			if cacheDuration == 0 {
				cacheDuration = 1 * time.Minute
			}
			Cache.Set(cacheKey, cacheJsonData, cacheDuration)
		}
	}
}