import time
import threading

class Snowflake:
    """
    雪花算法生成器：生成64位全局唯一ID
    结构：1位符号位（0） + 41位时间戳 + 10位机器ID + 12位序列号
    """
    def __init__(self, machine_id: int, start_time: float = 1577808000000):
        """
        初始化雪花算法生成器
        :param machine_id: 机器ID（0-1023，10位最大支持1024台机器）
        :param start_time: 起始时间戳（毫秒），默认2020-01-01 00:00:00的毫秒级时间戳
        """
        # 校验机器ID合法性（10位最大为2^10-1=1023）
        if machine_id < 0 or machine_id > 1023:
            raise ValueError("机器ID必须在0-1023之间")

        self.machine_id = machine_id  # 机器ID
        self.start_time = start_time  # 起始时间戳（毫秒）

        # 位偏移量定义
        self.timestamp_bits = 41  # 时间戳位数
        self.machine_bits = 10    # 机器ID位数
        self.sequence_bits = 12   # 序列号位数

        # 最大取值计算（位运算：n位最大为2^n -1）
        self.max_sequence = (1 << self.sequence_bits) - 1  # 4095
        self.max_timestamp = (1 << self.timestamp_bits) - 1  # 2^41-1

        # 偏移量（用于拼接各部分）
        self.machine_shift = self.sequence_bits  # 机器ID需要左移12位
        self.timestamp_shift = self.machine_bits + self.sequence_bits  # 时间戳需要左移22位

        # 状态变量（记录上一次生成ID的信息）
        self.last_timestamp = 0  # 上一次时间戳
        self.sequence = 0  # 当前毫秒内的序列号

        # 线程锁（保证多线程环境下序列号安全）
        self.lock = threading.Lock()

    def _get_current_timestamp(self) -> int:
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)

    def generate_id(self) -> int:
        """生成一个雪花ID"""
        with self.lock:  # 加锁保证线程安全
            current_timestamp = self._get_current_timestamp()

            # 1. 处理时钟回拨（当前时间 < 上一次时间）
            if current_timestamp < self.last_timestamp:
                # 计算回拨毫秒数，这里选择等待时间追平（也可抛出异常）
                sleep_ms = self.last_timestamp - current_timestamp
                time.sleep(sleep_ms / 1000)  # 转换为秒
                current_timestamp = self._get_current_timestamp()  # 重新获取时间

                # 若仍回拨，说明时间异常，抛出错误
                if current_timestamp < self.last_timestamp:
                    raise RuntimeError(f"时钟回拨异常，回拨{self.last_timestamp - current_timestamp}毫秒")

            # 2. 处理同一毫秒内的序列号
            if current_timestamp == self.last_timestamp:
                self.sequence += 1
                # 序列号超出最大值（4095），等待到下一个毫秒
                if self.sequence > self.max_sequence:
                    # 循环等待至下一毫秒
                    while current_timestamp <= self.last_timestamp:
                        current_timestamp = self._get_current_timestamp()
                    self.sequence = 0  # 重置序列号
            else:
                # 跨毫秒，重置序列号
                self.sequence = 0

            # 3. 更新上一次时间戳
            self.last_timestamp = current_timestamp

            # 4. 计算相对时间戳（当前时间 - 起始时间）
            relative_timestamp = current_timestamp - self.start_time
            if relative_timestamp > self.max_timestamp:
                raise RuntimeError("雪花算法已过期（超过41位时间戳上限）")

            # 5. 拼接ID：时间戳 << 22 | 机器ID << 12 | 序列号
            snowflake_id = (
                (relative_timestamp << self.timestamp_shift)
                | (self.machine_id << self.machine_shift)
                | self.sequence
            )
            return snowflake_id

if __name__ == "__main__":
    # 初始化生成器（机器ID=1，使用默认起始时间）
    snowflake = Snowflake(machine_id=1)

    # 生成10个ID并打印
    for _ in range(10):
        sid = snowflake.generate_id()
        print(f"雪花ID: {sid}（长度：{len(str(sid))}位）")

    # 多线程测试（验证线程安全）
    def generate_in_thread():
        for _ in range(5):
            print(f"线程{threading.current_thread().name}生成ID: {snowflake.generate_id()}")

    threads = [threading.Thread(target=generate_in_thread, name=f"thread-{i}") for i in range(3)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()