package main

import (
	"context"
	"fmt"
	"log"
	"net/http"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/mux"
)

var ctx = context.Background()

// 初始化 Redis 客户端
func initRedisClient() *redis.Client {
	client := redis.NewClient(&redis.Options{
		Addr: "localhost:6379", // Redis 地址
		DB:   3,                // 使用默认数据库
	})
	return client
}

// 限流的 Lua 脚本
var luaScript = `
  -- 定义返回值，是个数组，包含：是否触发限流（1限流 0通过）、当前桶中的令牌数
  local ret = {0, 0}
  local cl_key = KEYS[1]
  local lock_key = cl_key .. '-lock'
  local lock_val = redis.call('get', lock_key)
  if lock_val == '1' then
    ret[1] = 1
    ret[2] = -1
    return ret
  end

  -- 获取桶容量和其他配置
  local capacity = tonumber(ARGV[1])
  local amount = tonumber(ARGV[2])
  local inflow_unit = tonumber(ARGV[3])
  local inflow_quantity_per_unit = tonumber(ARGV[4])
  local key_expire_time = tonumber(ARGV[5])
  local lock_seconds = tonumber(ARGV[6])
  local current_time = tonumber(redis.call('time')[1])
  local st_key = cl_key .. '-st'

  -- 获取[上次向桶中投放令牌的时间]，如果没有设置过这个投放时间，则令牌桶也不存在
  local last_time = redis.call('get', st_key)
  if not last_time then
    local bucket_amount = capacity - amount
    redis.call('set', KEYS[1], bucket_amount, 'PX', key_expire_time)
    redis.call('set', st_key, current_time, 'PX', key_expire_time)
    ret[2] = bucket_amount
    return ret
  end

  -- 令牌桶存在，获取令牌桶中的当前令牌数
  local current_value = tonumber(redis.call('get', KEYS[1]))

  -- 判断是否需要投放令牌：当前时间 - 上次投放时间 >= 投放时间间隔
  local past_time = current_time - tonumber(last_time)
  local past_inflow_unit_quantity = math.floor(past_time / inflow_unit)

  local bucket_amount = current_value + past_inflow_unit_quantity * inflow_quantity_per_unit - amount
  -- 确保 bucket_amount 不会超过容量
  bucket_amount = math.min(bucket_amount, capacity - amount)

  local last_time_changed = 0

  if past_inflow_unit_quantity > 0 then
    last_time = current_time - (past_time % inflow_unit)
    redis.call('set', st_key, last_time, 'PX', key_expire_time)
    last_time_changed = 1
  end

  if bucket_amount < 0 then
    if lock_seconds > 0 then
      redis.call('set', lock_key, '1', 'EX', lock_seconds, 'NX')
    end
    ret[1] = 1
    ret[2] = -1
  else
    redis.call('set', KEYS[1], bucket_amount, 'PX', key_expire_time)
    ret[2] = bucket_amount
  end

  return ret
`

// 限流 API 处理函数
func rateLimitHandler(w http.ResponseWriter, r *http.Request) {
	client := initRedisClient()
	defer client.Close()

	// 请求参数
	keys := []string{"my-rate-limit-key"}
	argv := []interface{}{
		10,    // 桶容量
		1,     // 每次消耗的令牌数
		5,     // 投放令牌的时间间隔（秒）
		2,     // 每次投放的令牌数
		60000, // 键的过期时间（毫秒）
		0,     // 限流惩罚锁定时间（秒）
	}

	// 执行 Lua 脚本
	result, err := client.Eval(ctx, luaScript, keys, argv...).Result()
	if err != nil {
		http.Error(w, fmt.Sprintf("Error executing Lua script: %v", err), http.StatusInternalServerError)
		return
	}

	// 处理返回的结果
	ret := result.([]interface{})
	if ret[0].(int64) == 1 {
		// 触发限流
		http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
		return
	}

	// 返回当前桶中的令牌数
	_, _ = fmt.Fprintf(w, "Tokens left: %v", ret[1].(int64))
}

func main() {
	// 创建路由
	r := mux.NewRouter()

	// 注册 API 路由
	r.HandleFunc("/rate-limit", rateLimitHandler).Methods("GET")

	// 启动 HTTP 服务器
	log.Println("Server started on :8080")
	log.Fatal(http.ListenAndServe(":8080", r))
}