import threading
import time


class SnowflakeIDGenerator:
    # 目标：总比特数约40位，十进制最多12位（2^40 ≈ 1.1e12）
    # 时间戳用28位（可覆盖约8.5年），机器ID用6位（64台），序列号用6位（每毫秒64个）
    TIMESTAMP_BITS = 28
    WORKER_ID_BITS = 6
    SEQUENCE_BITS = 6

    # 起始时间戳（2025-01-01 00:00:00 UTC），单位毫秒
    CUSTOM_EPOCH = 1735689600000

    MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)
    MAX_SEQUENCE = -1 ^ (-1 << SEQUENCE_BITS)

    # 移位计算
    WORKER_ID_SHIFT = SEQUENCE_BITS
    TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS

    def __init__(self, worker_id, sequence=0):
        if worker_id > self.MAX_WORKER_ID or worker_id < 0:
            raise ValueError(f"Worker ID must be between 0 and {self.MAX_WORKER_ID}")
        self.worker_id = worker_id
        self.sequence = sequence
        self.last_timestamp = -1
        self.lock = threading.Lock()

    def _wait_next_ms(self, last_ts):
        ts = self._now()
        while ts <= last_ts:
            ts = self._now()
        return ts

    def _now(self):
        return int(time.time() * 1000)

    def generate_id(self):
        with self.lock:
            ts = self._now()
            if ts < self.last_timestamp:
                raise Exception("Clock moved backwards")
            if ts == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                if self.sequence == 0:
                    ts = self._wait_next_ms(self.last_timestamp)
            else:
                self.sequence = 0
            self.last_timestamp = ts

            # 相对时间（毫秒），只取低TIMESTAMP_BITS位
            relative_ts = (ts - self.CUSTOM_EPOCH) & ((1 << self.TIMESTAMP_BITS) - 1)
            new_id = (relative_ts << self.TIMESTAMP_SHIFT) | \
                     (self.worker_id << self.WORKER_ID_SHIFT) | \
                     self.sequence
            return new_id

    def generate_str(self):
        return str(self.generate_id())


# 全局生成器实例 (worker_id和datacenter_id需在分布式环境中保证唯一)
# 单机开发环境可暂时固定，生产环境建议通过配置中心或启动参数传入
generator = SnowflakeIDGenerator(worker_id=1)

if __name__ == "__main__":
    print(generator.generate_id())
    print(type(generator.generate_id()))
    print(generator.generate_str())