package ratelimit

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/patrickmn/go-cache"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	"golang.org/x/time/rate"

	"github.com/pingcap-inc/clinic/pkg/auth/session"
	"github.com/pingcap-inc/clinic/pkg/config"
)

type RateLimiter struct {
	store *cache.Cache
	sync.RWMutex

	limitCfgs []*config.LimitConfig
}

func NewRateLimiter(limitCfgs []*config.LimitConfig) *RateLimiter {
	return &RateLimiter{
		store:     cache.New(cache.NoExpiration, cache.NoExpiration),
		RWMutex:   sync.RWMutex{},
		limitCfgs: limitCfgs,
	}
}

func (s *RateLimiter) Get(key string) *limiter {
	v, ok := s.store.Get(key)
	if ok {
		return v.(*limiter)
	}

	s.Lock()
	defer s.Unlock()

	v, ok = s.store.Get(key)
	if ok {
		return v.(*limiter)
	}

	limiter := newLimiter(key, s.limitCfgs)
	s.store.Set(key, limiter, cache.NoExpiration)

	return limiter
}

type limiter struct {
	limiter *rate.Limiter
}

func (l *limiter) Allow() bool {
	return l.limiter.Allow()
}

func newLimiter(key string, sortLimits []*config.LimitConfig) *limiter {
	return match(key, sortLimits)
}

func match(key string, sortLimits []*config.LimitConfig) *limiter {
	for _, limit := range sortLimits {
		if URIMatch(key, limit.URI) {
			return &limiter{
				limiter: rate.NewLimiter(rate.Limit(limit.Limit), limit.Burst),
			}
		}
	}

	defaultLimit := sortLimits[len(sortLimits)-1]
	return &limiter{
		limiter: rate.NewLimiter(rate.Limit(defaultLimit.Limit), defaultLimit.Burst),
	}
}

func URIMatch(key1, key2 string) bool {
	key2 = strings.Replace(key2, "/*", "/.*", -1)

	re := regexp.MustCompile(`:[^/]+`)
	key2 = re.ReplaceAllString(key2, "$1[^/]+$2")

	return RegexMatch(key1, "^"+key2+"$")
}

func RegexMatch(key1, key2 string) bool {
	res, err := regexp.MatchString(key2, key1)
	if err != nil {
		log.Error("regex match failed", zap.String("key1", key1), zap.String("key2", key2))
		return false
	}
	return res
}

type RateLimiterMiddleware struct {
	limiter *RateLimiter
}

func NewRateLimiterMiddleware(cfg []*config.LimitConfig) *RateLimiterMiddleware {
	if len(cfg) == 0 {
		return &RateLimiterMiddleware{}
	}

	return &RateLimiterMiddleware{
		limiter: NewRateLimiter(cfg),
	}
}

func (rlm *RateLimiterMiddleware) Middleware(c *gin.Context) {
	if rlm.limiter == nil {
		c.Next()

		return
	}

	for _, uri := range session.IgnoreUris {
		if strings.HasPrefix(c.Request.RequestURI, uri) {
			c.Next()
			return
		}
	}

	user := session.GetUser(c)
	if user == nil {
		RateLimitError(c)
		return
	}

	key := fmt.Sprintf("%s/%d", c.Request.URL.Path, user.UserID)

	limiter := rlm.limiter.Get(key)
	if !limiter.Allow() {
		RateLimitError(c)

		return
	}

	c.Next()
}

func RateLimitError(c *gin.Context) {
	resp := map[string]string{"error": "too many requests"}

	c.JSON(http.StatusTooManyRequests, resp)
	c.Abort()
}