package templatecache

import (
	"bytes"
	"net/http"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/patrickmn/go-cache"
)

const (
	maxCacheItems = 2000
	maxCacheBytes = 64 << 20
)

var (
	Cache      *cache.Cache
	cacheBytes int64
)

// 章节详情、资讯详情等 ID 空间过大，禁止进进程缓存。
var highCardinalityContains = []string{
	"/chapter/info",
	"/consultinginfo/",
}

func init() {
	Cache = cache.New(5*time.Minute, 5*time.Minute)
	Cache.OnEvicted(func(_ string, v interface{}) {
		if b, ok := v.([]byte); ok {
			atomic.AddInt64(&cacheBytes, -int64(len(b)))
		}
	})
}

func shouldBypassCache(path string) bool {
	if path == "" {
		return false
	}
	if i := strings.IndexByte(path, '?'); i >= 0 {
		path = path[:i]
	}
	for _, part := range highCardinalityContains {
		if strings.Contains(path, part) {
			return true
		}
	}
	return false
}

func cacheBudgetAllows(n int) bool {
	if n <= 0 {
		return false
	}
	if Cache.ItemCount() >= maxCacheItems {
		return false
	}
	return atomic.LoadInt64(&cacheBytes)+int64(n) <= maxCacheBytes
}

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) {
		if ctx.Request.Method != "GET" {
			ctx.Next()
			return
		}
		if shouldBypassCache(ctx.Request.URL.Path) {
			ctx.Next()
			return
		}

		cacheKey := ctx.Request.URL.String()
		if cacheData, exists := Cache.Get(cacheKey); exists {
			if item, ok := cacheData.([]byte); ok && len(item) > 0 {
				ctx.Data(http.StatusOK, "application/json; charset=utf-8", item)
				ctx.Abort()
				return
			}
		}

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

		if ctx.Writer.Status() != 200 {
			return
		}
		raw := writer.body.Bytes()
		if !cacheBudgetAllows(len(raw)) {
			return
		}
		copied := append([]byte(nil), raw...)
		atomic.AddInt64(&cacheBytes, int64(len(copied)))
		cacheDuration := duration
		if cacheDuration == 0 {
			cacheDuration = time.Minute
		}
		Cache.Set(cacheKey, copied, cacheDuration)
	}
}