package generator

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

const (
	epoch            = int64(1609459200000)
	datacenterIdBits = uint(5)
	workerIdBits     = uint(5)
	sequenceBits     = uint(12)

	maxDatacenterId = -1 ^ (-1 << datacenterIdBits)
	maxWorkerId     = -1 ^ (-1 << workerIdBits)
	maxSequence     = -1 ^ (-1 << sequenceBits)

	workerIdShift      = sequenceBits
	datacenterIdShift  = sequenceBits + workerIdBits
	timestampLeftShift = sequenceBits + workerIdBits + datacenterIdBits
)

type Snowflake struct {
	mu           sync.Mutex
	timestamp    int64
	datacenterId int64
	workerId     int64
	sequence     int64
}

func NewSnowflake(datacenterId, workerId int64) (*Snowflake, error) {
	if datacenterId < 0 || datacenterId > maxDatacenterId {
		return nil, errors.New("datacenter id must be between 0 and 31")
	}
	if workerId < 0 || workerId > maxWorkerId {
		return nil, errors.New("worker id must be between 0 and 31")
	}

	return &Snowflake{
		timestamp:    0,
		datacenterId: datacenterId,
		workerId:     workerId,
		sequence:     0,
	}, nil
}

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

	now := time.Now().UnixMilli()

	if now < s.timestamp {
		return 0, errors.New("clock moved backwards")
	}

	if now == s.timestamp {
		s.sequence = (s.sequence + 1) & maxSequence
		if s.sequence == 0 {
			for now <= s.timestamp {
				now = time.Now().UnixMilli()
			}
		}
	} else {
		s.sequence = 0
	}

	s.timestamp = now

	id := ((now - epoch) << timestampLeftShift) |
		(s.datacenterId << datacenterIdShift) |
		(s.workerId << workerIdShift) |
		s.sequence

	return id, nil
}

func (s *Snowflake) NextIDs(count int) ([]int64, error) {
	if count <= 0 || count > 1000 {
		return nil, errors.New("count must be between 1 and 1000")
	}

	ids := make([]int64, count)
	for i := 0; i < count; i++ {
		id, err := s.NextID()
		if err != nil {
			return nil, err
		}
		ids[i] = id
	}

	return ids, nil
}