package com.sonic.common.ratelimit.impl;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import com.sonic.common.AppRuntime;
import com.sonic.common.ratelimit.RateLimiter;
import org.springframework.data.redis.core.*;

import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.hash.Hashing;

import lombok.Builder;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class RedisRateLimiter implements RateLimiter {
    private AppRuntime appRuntime;
    private RedisTemplate redisTemplate;
    private RateLimiterWorker rateLimiterWorker;
    /**
     * 自定义的key, 避免redisKey冲突. 必填
     */
    private String limiterKey;
    private List<LimitRule> limitRules;
    /**
     * 窗口保存最大时长. 主要针对Sliding方式窗口的zSet过期
     */
    private Integer maxWindowDurationHour;
    private WindowType windowType;

    @Builder
    RedisRateLimiter(AppRuntime appRuntime,
                     RedisTemplate redisTemplate,
                     WindowType windowType,
                     String limiterKey,
                     List<LimitRule> limitRules,
                     Integer maxWindowDurationHour) {
        Objects.requireNonNull(appRuntime);
        Objects.requireNonNull(redisTemplate);
        Objects.requireNonNull(windowType);
        Objects.requireNonNull(limitRules);
        Objects.requireNonNull(limiterKey);
        Preconditions.checkArgument(!limitRules.isEmpty());
        if (limitRules.stream().anyMatch(p -> !p.isValid())) {
            throw new RuntimeException(String.format("invalid rate expression, expression = %s", JSONObject.toJSONString(limitRules)));
        }

        this.windowType = windowType;
        this.appRuntime = appRuntime;
        this.redisTemplate = redisTemplate;
        this.limitRules = limitRules;
        this.limiterKey = limiterKey;

        this.maxWindowDurationHour = Optional.ofNullable(maxWindowDurationHour).orElse(24);

        this.rateLimiterWorker = buildWorker(windowType);
    }

    @Override
    public Optional<Permit> tryAcquire(Object value) {
        if (!rateLimiterWorker.tryAcquire(value)) {
            return Optional.empty();
        }
        List<Runnable> cancelRunners = rateLimiterWorker.visit(value);
        return Optional.of(Permit.builder()
                .cancelRunners(cancelRunners)
                .build());
    }

    @Override
    public WindowType getWindowType() {
        return windowType;
    }

    private RateLimiterWorker buildWorker(WindowType windowType) {
        if (windowType == WindowType.FIXED) {
            return new FixedWindowRateLimiter();
        }
        if (windowType == WindowType.FIXED_BUCKET) {
            return new FixedBucketWindowRateLimiter();
        }
        if (windowType == WindowType.SLIDING) {
            return new SlidingWindowRateLimiter();
        }
        throw new RuntimeException(String.format("unsupported window type, window type = %s", windowType));
    }

    private String buildRedisKey(Object value) {
        String hash = Hashing.murmur3_128().newHasher()
                .putString(limiterKey, Charsets.UTF_8)
                .putString(String.valueOf(value), Charsets.UTF_8)
                .hash()
                .toString();
        return appRuntime.buildPrefixKey("rl", getWindowType().getShortName(), hash);
    }

    /**
     * 固定窗口限流, 窗口起始时间第一次tryAcquire时时间. 窗口大小为WindowDuration
     * <pre>
     *     key = value + WindowDuration.
     *     在该窗口被访问时, 计数器+1. 窗口持续时长为WindowDuration. 并依赖redis ttl销毁
     *     窗口被销毁后, 重置计数器
     * </pre>
     */
    class FixedWindowRateLimiter implements RateLimiterWorker {
        @Override
        public List<Runnable> visit(Object value) {
            List<Runnable> cancels = new ArrayList<>(limitRules.size());
            //根据时间构建key, 避免hash无法删除
            limitRules.stream()
                    .forEach(e -> {
                        String key = buildLimiterKey(value, e);
                        BoundValueOperations op = redisTemplate.boundValueOps(key);
                        Long result = op.increment(1);
                        //第一次访问时设置过期时间, 以确定窗口
                        if (result == 1) {
                            redisTemplate.expire(key, e.getSeconds(), TimeUnit.SECONDS);
                        }
                        cancels.add(() -> op.decrement(1));
                    });
            return cancels;
        }

        @Override
        public boolean tryAcquire(Object value) {
            boolean anyMatch = limitRules.stream()
                    .anyMatch(limitRule -> {
                        String key = buildLimiterKey(value, limitRule);
                        return limitRule.getPermits() <= Optional.ofNullable(redisTemplate.opsForValue().get(key))
                                .map(e -> Long.parseLong(e.toString()))
                                .orElse(0L);
                    });
            return !anyMatch;
        }

        protected String buildLimiterKey(Object value, LimitRule limitRule) {
            return buildRedisKey(value + ":" + limitRule.getSeconds());
        }
    }

    /**
     * 滑动窗口限流, 每次获取令牌成功时间加入到zset. 后续获取令牌时每次检查zset中WindowDuration中已获取令牌数. 并判断是否可以继续获取令牌
     * <pre>
     *     key = value
     *     zset value = currentMillis. score = currentMillis
     *     获取令牌时, 在计算zset中 score = [currentMillis-WindowDuration, currentMillis} 的element数量
     * </pre>
     */
    class SlidingWindowRateLimiter implements RateLimiterWorker {
        //当zset的element达到一定数量时, 清理该zet. 避免redis内存泄露
        private static final int CLEAN_KEY_THRESHOLD = 1000;
        private AtomicLong visitCounter = new AtomicLong();

        @Override
        public List<Runnable> visit(Object value) {
            String key = buildRedisKey(value);
            long now = System.currentTimeMillis();
            String member = String.valueOf(now);
            final BoundZSetOperations op = redisTemplate.boundZSetOps(key);
            op.add(member, now);

            redisTemplate.expire(key, maxWindowDurationHour, TimeUnit.HOURS);
            if (visitCounter.incrementAndGet() > CLEAN_KEY_THRESHOLD) {
                //删除过期的访问记录
                op.removeRangeByScore(0, now - TimeUnit.HOURS.toMillis(maxWindowDurationHour));
                visitCounter.set(0);
            }
            return ImmutableList.of(() -> op.remove(member));
        }

        @Override
        public boolean tryAcquire(Object value) {
            String key = buildRedisKey(value);
            long now = System.currentTimeMillis();

            //检查所有的rule, 如果有其中一个失败则跳出
            return !limitRules.stream()
                    .anyMatch(p -> p.getPermits() <= redisTemplate.opsForZSet().count(key, now - TimeUnit.SECONDS.toMillis(p.getSeconds()), now));
        }
    }

    /**
     * 固定窗口限流, 窗口根据自然时间向前滚动
     * 继承自FixedWindowRateLimiter, 区别在于构建key的方式不一样
     * <pre>
     *     key = value + currentMillis/WindowDuration.
     *     currentMillis/WindowDuration会把自然时间分割为长为WindowDuration的片段. 片段有效期为WindowDuration
     *     获取令牌时在该片段上检查是否有剩余令牌
     * </pre>
     */
    class FixedBucketWindowRateLimiter extends FixedWindowRateLimiter implements RateLimiterWorker {
        @Override
        protected String buildLimiterKey(Object value, LimitRule limitRule) {
            return buildRedisKey(value + ":" + System.currentTimeMillis() / TimeUnit.SECONDS.toMillis(limitRule.getSeconds()));
        }
    }

    interface RateLimiterWorker {
        /**
         * 尝试获取令牌
         *
         * @param value
         * @return 如果获取成功则返回true, 失败则为false
         */
        boolean tryAcquire(Object value);

        /**
         * 获取令牌完成后增加计数器
         *
         * @param value
         * @return limit key
         */
        List<Runnable> visit(Object value);
    }
}
