package main

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

const (
	// 位数分配 (Bit Allocation)
	timestampBits  uint8 = 41
	businessIDBits uint8 = 3
	workerIDBits   uint8 = 10
	sequenceBits   uint8 = 9 // 64 - 1 (sign) - 41 (timestamp) - 3 (business) - 10 (worker) = 9

	// 最大值 (Max Values)
	maxBusinessID int64 = -1 ^ (-1 << businessIDBits) // 2^3 - 1 = 7
	maxWorkerID   int64 = -1 ^ (-1 << workerIDBits)   // 2^10 - 1 = 1023
	maxSequence   int64 = -1 ^ (-1 << sequenceBits)   // 2^9 - 1 = 511

	// 位移量 (Bit Shifts)
	workerIDShift   = sequenceBits                                 // 9
	businessIDShift = sequenceBits + workerIDBits                  // 9 + 10 = 19
	timestampShift  = sequenceBits + workerIDBits + businessIDBits // 9 + 10 + 3 = 22

	// 自定义纪元 (Epoch), 单位毫秒. (可以设置为项目上线的日期)
	// 例如: 2025-01-01 00:00:00 UTC
	customEpoch int64 = 1735689600000 // time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC).UnixMilli()
)

// ErrClockMovedBackwards indicates that the system clock moved backwards.
var ErrClockMovedBackwards = errors.New("clock moved backwards, refusing to generate id")

// ErrInvalidBusinessID indicates the business ID is out of range.
var ErrInvalidBusinessID = fmt.Errorf("business ID must be between 0 and %d", maxBusinessID)

// ErrInvalidWorkerID indicates the worker ID is out of range.
var ErrInvalidWorkerID = fmt.Errorf("worker ID must be between 0 and %d", maxWorkerID)

// Generator is the core ID generator structure.
type Generator struct {
	mu            sync.Mutex
	lastTimestamp int64
	workerID      int64
	businessID    int64
	sequence      int64
}

// NewGenerator creates a new ID generator instance.
func NewGenerator(workerID, businessID int64) (*Generator, error) {
	if workerID < 0 || workerID > maxWorkerID {
		return nil, ErrInvalidWorkerID
	}
	if businessID < 0 || businessID > maxBusinessID {
		return nil, ErrInvalidBusinessID
	}

	return &Generator{
		workerID:   workerID,
		businessID: businessID,
	}, nil
}

// currentTimeMillis returns the current time in milliseconds since the custom epoch.
func currentTimeMillis() int64 {
	return time.Now().UnixMilli() - customEpoch
}

// tilNextMillis waits until the next millisecond.
func tilNextMillis(lastTimestamp int64) int64 {
	timestamp := currentTimeMillis()
	for timestamp <= lastTimestamp {
		// Spin wait is generally discouraged, but for millisecond precision,
		// time.Sleep(1 * time.Millisecond) might overshoot too much.
		// A brief sleep can reduce CPU churn if clock skew is minor.
		time.Sleep(time.Microsecond * 100) // Sleep briefly
		timestamp = currentTimeMillis()
	}
	return timestamp
}

// NextID generates the next unique ID.
func (g *Generator) NextID() (int64, error) {
	g.mu.Lock()
	defer g.mu.Unlock()

	timestamp := currentTimeMillis()

	// 时钟回拨检查 (Clock moved backwards check)
	if timestamp < g.lastTimestamp {
		// 可以选择:
		// 1. 返回错误 (Recommended for safety)
		return 0, fmt.Errorf("%w: current: %d, last: %d", ErrClockMovedBackwards, timestamp, g.lastTimestamp)
		// 2. 等待时钟追上 (Potentially blocks, less safe if clock jump is large)
		// timestamp = tilNextMillis(g.lastTimestamp)
	}

	// 同一毫秒内 (Within the same millisecond)
	if timestamp == g.lastTimestamp {
		g.sequence = (g.sequence + 1) & maxSequence
		// 序列号溢出，等待下一毫秒 (Sequence overflow, wait for next millisecond)
		if g.sequence == 0 {
			timestamp = tilNextMillis(g.lastTimestamp)
			// 等待后重置序列号 (Reset sequence after waiting)
			// g.sequence = 0 // Reset is implicit as it overflowed to 0
		}
	} else {
		// 新的毫秒，重置序列号 (New millisecond, reset sequence)
		g.sequence = 0
	}

	// 更新最后时间戳 (Update last timestamp)
	g.lastTimestamp = timestamp

	// 组合 ID (Assemble the ID)
	id := (timestamp << timestampShift) | // 时间戳左移
		(g.businessID << businessIDShift) | // 业务ID左移
		(g.workerID << workerIDShift) | // Worker ID左移
		g.sequence // 序列号

	return id, nil
}

// ParseID decomposes an ID into its components. Useful for debugging or analysis.
func ParseID(id int64) (timestampMsSinceEpoch int64, businessID int64, workerID int64, sequence int64, genTimeUTC time.Time) {
	timestampMsSinceEpoch = (id >> timestampShift) & (int64(-1) ^ (int64(-1) << timestampBits))
	businessID = (id >> businessIDShift) & maxBusinessID
	workerID = (id >> workerIDShift) & maxWorkerID
	sequence = id & maxSequence

	// Calculate generation time in UTC
	genTimeUTC = time.UnixMilli(timestampMsSinceEpoch + customEpoch).UTC()

	return
}

func main() {
	// !!! IMPORTANT: Worker ID and Business ID MUST be unique per instance/purpose !!!
	// These should typically come from configuration.
	workerID := int64(1)          // Example: Get from config/env
	businessIDOrder := int64(1)   // Example: 1 for Order
	businessIDPayment := int64(2) // Example: 2 for Payment

	// Create generators for different business types
	orderGenerator, err := NewGenerator(workerID, businessIDOrder)
	if err != nil {
		panic(err)
	}

	paymentGenerator, err := NewGenerator(workerID, businessIDPayment)
	if err != nil {
		panic(err)
	}

	// Generate some IDs
	for i := 0; i < 5; i++ {
		orderID, err := orderGenerator.NextID()
		if err != nil {
			fmt.Println("Error generating order ID:", err)
			time.Sleep(1 * time.Millisecond) // Wait before retry on clock issues
			continue
		}
		fmt.Printf("Generated Order ID: %d\n", orderID)
		// You can optionally add a human-readable prefix when displaying/logging
		fmt.Printf("   Readable Order ID: ORD-%d\n", orderID)

		// Parse it back (for demonstration)
		ts, biz, wkr, seq, genTime := ParseID(orderID)
		fmt.Printf("   Parsed: Timestamp=%d, Business=%d, Worker=%d, Sequence=%d, GenTime=%s\n", ts, biz, wkr, seq, genTime.Format(time.RFC3339Nano))

		paymentID, err := paymentGenerator.NextID()
		if err != nil {
			fmt.Println("Error generating payment ID:", err)
			continue
		}
		fmt.Printf("Generated Payment ID: %d\n", paymentID)
		fmt.Printf("   Readable Payment ID: PAY-%d\n", paymentID)

		tsP, bizP, wkrP, seqP, genTimeP := ParseID(paymentID)
		fmt.Printf("   Parsed: Timestamp=%d, Business=%d, Worker=%d, Sequence=%d, GenTime=%s\n\n", tsP, bizP, wkrP, seqP, genTimeP.Format(time.RFC3339Nano))

		time.Sleep(5 * time.Millisecond) // Simulate time passing
	}

	// Example of how to get shard key (assuming 64 shards)
	orderIDForSharding, _ := orderGenerator.NextID()
	shardIndex := orderIDForSharding % 64
	fmt.Printf("\nOrder ID %d would route to Shard Index %d\n", orderIDForSharding, shardIndex)
}