import os
import time

# 自定义 Base62 字符表
ALPHABET = "Kvjk3l5GT8Z4xLMSiDdW1nfE9az0HbUgoJ7cQ6pqBmXCwYFReh2NyIOrAPVstu"

MACHINE_ID = os.getenv("SHORT_SNOWFLAKE_MACHINE_ID")
MACHINE_ID = int(MACHINE_ID) if MACHINE_ID is not None else 29


def encode_base62(num: int) -> str:
    if num == 0:
        return "0"
    s = []
    base = 62
    while num > 0:
        num, r = divmod(num, base)
        s.append(ALPHABET[r])
    return "".join(reversed(s))


class ShortSnowflake:
    """
    Python 版短 ID Snowflake
    结构：41 bits timestamp | 10 bits machine | 12 bits counter = 63 bits
    """

    def __init__(self, machine_id: int = 1):
        self.machine_id = machine_id & ((1 << 10) - 1)  # 10 bits
        self.counter = 0
        self.last_ts = 0

    def next(self) -> str:
        ts = int(time.time() * 1000)  # 毫秒

        if ts == self.last_ts:
            # 同毫秒递增计数
            self.counter = (self.counter + 1) & ((1 << 12) - 1)
            if self.counter == 0:
                # 计数溢出，等待下一毫秒
                while int(time.time() * 1000) == ts:
                    pass
                ts = int(time.time() * 1000)
        else:
            self.counter = 0
            self.last_ts = ts

        # 组合 Snowflake
        id_num = ((ts & ((1 << 41) - 1)) << 22) | (self.machine_id << 12) | self.counter

        # 任意大小 int → base62
        return encode_base62(id_num)


def new_id(machine_id: int = MACHINE_ID, prefix: str = "") -> str:
    """生成新的短 ID。"""
    global _short_snowflake_instance
    if "_short_snowflake_instance" not in globals():
        _short_snowflake_instance = ShortSnowflake(machine_id=machine_id)
    return prefix + _short_snowflake_instance.next()
