package snowflake

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

const (
    // 起始时间戳 (2024-01-01 00:00:00 UTC)
    epoch int64 = 1704067200000

    // 各部分位数
    timestampBits  = 41 // 时间戳占用位数
    datacenterBits = 5  // 数据中心ID占用位数
    workerBits     = 5  // 工作机器ID占用位数
    sequenceBits   = 12 // 序列号占用位数

    // 最大值
    maxDatacenterID = -1 ^ (-1 << datacenterBits) // 31
    maxWorkerID     = -1 ^ (-1 << workerBits)     // 31
    maxSequence     = -1 ^ (-1 << sequenceBits)   // 4095

    // 位移量
    workerShift     = sequenceBits                         // 12
    datacenterShift = sequenceBits + workerBits            // 17
    timestampShift  = sequenceBits + workerBits + datacenterBits // 22
)

// Snowflake ID生成器
type Snowflake struct {
    mu           sync.Mutex // 互斥锁，保证并发安全
    timestamp    int64      // 上次生成ID的时间戳
    datacenterID int64      // 数据中心ID
    workerID     int64      // 工作机器ID
    sequence     int64      // 序列号
}

// 创建一个新的Snowflake实例
func NewSnowflake(datacenterID, workerID int64) (*Snowflake, error) {
    // 参数校验
    if datacenterID < 0 || datacenterID > maxDatacenterID {
        return nil, errors.New("datacenterID must be between 0 and 31")
    }
    if workerID < 0 || workerID > maxWorkerID {
        return nil, errors.New("workerID must be between 0 and 31")
    }

    return &Snowflake{
        timestamp:    0,
        datacenterID: datacenterID,
        workerID:     workerID,
        sequence:     0,
    }, nil
}

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

    // 获取当前时间戳(毫秒)
    now := time.Now().UnixMilli()

    // 如果当前时间小于上次生成ID的时间戳，说明发生了时钟回拨
    if now < s.timestamp {
        return 0, errors.New("clock moved backwards")
    }

    // 如果是同一毫秒内生成的ID
    if now == s.timestamp {
        // 序列号自增
        s.sequence = (s.sequence + 1) & maxSequence
        // 如果序列号溢出(超过4095)
        if s.sequence == 0 {
            // 阻塞到下一毫秒
            now = s.waitNextMillis(s.timestamp)
        }
    } else {
        // 不同毫秒，序列号重置为0
        s.sequence = 0
    }

    // 更新时间戳
    s.timestamp = now

    // 组装64位ID:
    // (时间戳 << 22) | (数据中心ID << 17) | (机器ID << 12) | 序列号
    id := ((now - epoch) << timestampShift) |
        (s.datacenterID << datacenterShift) |
        (s.workerID << workerShift) |
        s.sequence

    return id, nil
}

// 阻塞等待到下一毫秒
func (s *Snowflake) waitNextMillis(lastTimestamp int64) int64 {
    timestamp := time.Now().UnixMilli()
    for timestamp <= lastTimestamp {
        timestamp = time.Now().UnixMilli()
    }
    return timestamp
}

// 从ID中解析出各个组成部分
func ParseID(id int64) map[string]int64 {
    return map[string]int64{
        "timestamp":    ((id >> timestampShift) + epoch),
        "datacenterID": (id >> datacenterShift) & maxDatacenterID,
        "workerID":     (id >> workerShift) & maxWorkerID,
        "sequence":     id & maxSequence,
    }
}

// 从ID中提取生成时间
func GetTimestamp(id int64) time.Time {
    timestamp := (id >> timestampShift) + epoch
    return time.UnixMilli(timestamp)
}

// 使用Redis自动分配机器ID
func AutoAssignWorkerID(redis *redis.Client, datacenterID int64) (int64, error) {
    key := fmt.Sprintf("snowflake:datacenter:%d:workers", datacenterID)
    workerID, err := redis.Incr(ctx, key).Result()
    if err != nil {
        return 0, err
    }
    if workerID > maxWorkerID {
        return 0, errors.New("no available worker ID")
    }
    return workerID, nil
}