package main

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// Snowflake 雪花算法生成器
type Snowflake struct {
	mu        sync.Mutex // 保证并发安全
	epoch     int64      // 起始时间戳（毫秒）
	nodeID    int64      // 节点ID
	sequence  int64      // 毫秒内序列号
	lastStamp int64      // 上一次生成ID的时间戳（毫秒）

	// 位段掩码与偏移量（根据64位结构定义）
	nodeIDBits  int64 = 10                     // 节点ID位数
	sequenceBits int64 = 12                    // 序列号位数
	nodeIDMax   int64 = -1 ^ (-1 << nodeIDBits) // 节点ID最大值（1023）
	sequenceMax int64 = -1 ^ (-1 << sequenceBits) // 序列号最大值（4095）

	nodeIDShift  int64 = sequenceBits          // 节点ID偏移量
	timeStampShift int64 = sequenceBits + nodeIDBits // 时间戳偏移量
}

// NewSnowflake 创建雪花算法实例
// epoch: 起始时间戳（毫秒），如 1704067200000（2024-01-01 00:00:00）
// nodeID: 节点ID（0-1023）
func NewSnowflake(epoch int64, nodeID int64) (*Snowflake, error) {
	// 校验节点ID范围
	if nodeID < 0 || nodeID > (-1 ^ (-1 << 10)) {
		return nil, errors.New("node ID must be between 0 and 1023")
	}

	return &Snowflake{
		epoch:     epoch,
		nodeID:    nodeID,
		sequence:  0,
		lastStamp: -1,
	}, nil
}

// 获取当前毫秒级时间戳
func (s *Snowflake) now() int64 {
	return time.Now().UnixMilli()
}

// NextID 生成下一个唯一ID
func (s *Snowflake) NextID() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 1. 获取当前时间戳
	now := s.now()

	// 2. 处理时钟回拨
	if now < s.lastStamp {
		return 0, fmt.Errorf("clock moved backwards. Refusing to generate ID for %d milliseconds", s.lastStamp-now)
	}

	// 3. 同一毫秒内，序列号自增
	if now == s.lastStamp {
		s.sequence = (s.sequence + 1) & s.sequenceMax
		// 序列号溢出（同一毫秒生成超过4096个ID）
		if s.sequence == 0 {
			// 等待到下一毫秒
			for now <= s.lastStamp {
				now = s.now()
			}
		}
	} else {
		// 不同毫秒，重置序列号
		s.sequence = 0
	}

	// 4. 更新上一次时间戳
	s.lastStamp = now

	// 5. 组合64位ID
	// 时间戳部分：(now - epoch) << 22
	// 节点ID部分：nodeID << 12
	// 序列号部分：sequence
	id := ((now - s.epoch) << s.timeStampShift) | (s.nodeID << s.nodeIDShift) | s.sequence

	return id, nil
}

// 测试示例
func main() {
	// 自定义起始时间（2024-01-01 00:00:00）
	epoch := int64(1704067200000)
	// 节点ID（假设当前节点为1）
	nodeID := int64(1)

	// 创建雪花算法实例
	sf, err := NewSnowflake(epoch, nodeID)
	if err != nil {
		panic(err)
	}

	// 并发生成10个ID测试
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			id, err := sf.NextID()
			if err != nil {
				fmt.Printf("生成ID失败：%v\n", err)
				return
			}
			fmt.Printf("第%d个ID：%d\n", idx+1, id)
		}(i)
	}
	wg.Wait()
}