package cn.axzo.workflow.core.engine.id;

import cn.hutool.core.date.DateUtil;
import org.flowable.common.engine.impl.cfg.IdGenerator;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.ArrayList;
import java.util.List;

/**
 * 支持分布式的基于时间的 ID 生成器
 *
 * @author wangli
 * @since 2023/12/20 21:41
 */
public class DistributedTimeBasedIdGenerator implements IdGenerator {

    private final StringRedisTemplate stringRedisTemplate;

    public DistributedTimeBasedIdGenerator(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    private static final String LUA_SCRIPT = "local key = tostring(KEYS[1]);\n"
            + "local count = tonumber(KEYS[2]);\n"
            + "local dateStr = tostring(KEYS[3]);\n"
            + "\n"
            + "local newKey = key .. \"_\" .. dateStr;\n"
            + "local numRedis = redis.call(\"incr\", newKey);\n"
            + "print(numRedis);\n"
            + "\n"
            + "if (numRedis == 1) then\n"
            + "\tredis.call(\"expire\",newKey,60);\n"
            + "end\n"
            + "\n"
            + "local function DightNum(num)\n"
            + "\tif math.floor(num) ~= num or num < 0 then\n"
            + "\t\treturn -1;\n"
            + "\telseif 0 == num then\n"
            + "\t\treturn 1;\n"
            + "\telse\n"
            + "\t\tlocal tmp_dight = 0;\n"
            + "\t\twhile num > 0 do\n"
            + "\t\t\tnum = math.floor(num/10);\n"
            + "\t\t\ttmp_dight = tmp_dight + 1;\n"
            + "\t\tend\n"
            + "\t\treturn tmp_dight;\n"
            + "\tend\n"
            + "end\n"
            + "\n"
            + "local function AddZeroFrontNum(dest_dight, num)\n"
            + "\tlocal num_dight = DightNum(num);\n"
            + "\tif -1 == num_dight then\n"
            + "\t\treturn -1;\n"
            + "\telseif dest_dight <= num_dight then\n"
            + "\t\treturn tostring(num);\n"
            + "\telse\n"
            + "\t\tlocal str_e = \"\"\n"
            + "\t\tfor var =1, dest_dight - num_dight do\n"
            + "\t\t\tstr_e = str_e .. \"0\";\n"
            + "\t\tend\n"
            + "\t\treturn str_e .. tostring(num);\n"
            + "\tend\n"
            + "end\n"
            + "\n"
            + "local idStr = AddZeroFrontNum(count, numRedis);\n"
            + "return dateStr .. idStr;\n";

    @Override
    public synchronized String getNextId() {
        String now = DateUtil.format(DateUtil.date(), "yyyyMMddHHmm");

        DefaultRedisScript<String> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(LUA_SCRIPT);
        redisScript.setResultType(String.class);
        List<String> keys = new ArrayList<>();
        keys.add("Custom:Id:Generator");
        keys.add("9");
        keys.add(now);
        return stringRedisTemplate.execute(redisScript, keys);
    }
}
