import json
import time
from dataclasses import dataclass
from typing import List

from django.conf import settings

import promptgame.gameui.redis_utils as ru
from promptgame.settings.custom_types import RateLimitBucketConfig

LUA_SCRIPT = """
-- redis script fails if we don't have exactly 1 arg
if #ARGV ~= 1 then
    return redis.error_reply("ERR this script expects exactly one argument")
end

-- parse the args
local args = cjson.decode(ARGV[1])

local key_prefix = args.key_prefix
local current_time = args.current_time
local bucket_configs = args.bucket_configs

local min_wait_time = 0

for i, config in ipairs(bucket_configs) do
    config.last_update_time_key = key_prefix .. ":" .. config.bucket_name .. ":last_update_time"
    config.tokens_key = key_prefix .. ":" .. config.bucket_name .. ":tokens"

    config.last_update_time = tonumber(redis.call("GET", config.last_update_time_key)) or current_time
    config.tokens = tonumber(redis.call("GET", config.tokens_key)) or config.max_capacity

    -- Calculate the new number of tokens. math.max(..., 0) is to avoid negative
    -- time intervals, and math.min(tokens + delta, config.max_capacity) is to avoid
    -- overflowing the bucket.
    local delta = math.max((current_time - config.last_update_time) * config.refill_rate, 0)
    config.tokens = math.min(config.tokens + delta, config.max_capacity)

    -- Check if there's at least one token left
    if config.tokens < 1 then
        local wait_time = (1 - config.tokens) / config.refill_rate
        min_wait_time = math.max(min_wait_time, wait_time)
    end
end

if min_wait_time <= 0 then
    -- successful call, deduct tokens
    for i, config in ipairs(bucket_configs) do
        -- Remove a token and calculate time required for the bucket to fill
        config.tokens = config.tokens - 1
        local time_to_full = math.max((config.max_capacity - config.tokens) / config.refill_rate, 0)
        -- Expiry times must be whole numbers
        local expiration_time = math.ceil(time_to_full)
        redis.call("SET", config.tokens_key, config.tokens, "EX", expiration_time)
        redis.call("SET", config.last_update_time_key, current_time, "EX", expiration_time)
    end
end

-- return as string, otherwise Redis coerces the result to int
return tostring(min_wait_time)
"""


@dataclass
class RateLimitResult:
    request_allowed: bool
    wait_time_s: float


async def rate_limit_check(
    account_id: int,
    bucket_configs: List[RateLimitBucketConfig],
) -> RateLimitResult:
    client = await ru.get_async_client()
    script = client.register_script(LUA_SCRIPT)

    current_time = time.time()
    key_prefix = f"{settings.REDIS_KEY_PREFIX}:ratelim:{account_id}"
    lua_args = {
        "key_prefix": key_prefix,
        "current_time": current_time,
        "bucket_configs": bucket_configs,
    }

    result = await script(args=[json.dumps(lua_args)])
    assert isinstance(result, bytes)
    min_wait_time = float(result.decode("utf8"))

    return RateLimitResult(
        request_allowed=min_wait_time <= 0,
        wait_time_s=min_wait_time,
    )
