package utils

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

const (
	epoch          = int64(1609459200000) // 2021-01-01 00:00:00 的时间戳（毫秒）
	workerBits     = 5
	sequenceBits   = 12
	maxWorkerID    = -1 ^ (-1 << workerBits)
	maxSequenceID  = -1 ^ (-1 << sequenceBits)
	workerShift    = sequenceBits
	timestampShift = sequenceBits + workerBits
)

type Snowflake struct {
	mu         sync.Mutex
	timestamp  int64
	workerID   int
	sequenceID int
}

func NewSnowflake(workerID int) *Snowflake {
	if workerID < 0 || workerID > maxWorkerID {
		panic(fmt.Sprintf("worker ID must be between 0 and %d", maxWorkerID))
	}
	return &Snowflake{
		timestamp:  0,
		workerID:   workerID,
		sequenceID: 0,
	}
}

func (s *Snowflake) NextID() int64 {
	s.mu.Lock()
	defer s.mu.Unlock()

	currentTime := time.Now().UnixNano() / 1000000 // 转换为毫秒
	if s.timestamp == currentTime {
		s.sequenceID = (s.sequenceID + 1) & maxSequenceID
		if s.sequenceID == 0 {
			// 当前毫秒内的序列号用尽，等待下一毫秒
			for currentTime <= s.timestamp {
				currentTime = time.Now().UnixNano() / 1000000
			}
		}
	} else {
		s.sequenceID = 0
	}

	s.timestamp = currentTime
	id := ((currentTime - epoch) << timestampShift) |
		(int64(s.workerID) << workerShift) |
		int64(s.sequenceID)
	return id
}