package id_generator

import (
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"net"

	"github.com/bwmarrin/snowflake"
)

type SnowflakeIDGenerator struct {
	node *snowflake.Node
}

func MustNewSnowflakeIDGenerator(machineIDGenerator MachineIDGenerator) IDGenerator {
	nodeID := machineIDGenerator.MustGenerateMachineID()
	node, err := snowflake.NewNode(int64(nodeID))
	if err != nil {
		panic(err)
	}
	return &SnowflakeIDGenerator{
		node: node,
	}
}

func (s *SnowflakeIDGenerator) NextID() uint64 {
	return uint64(s.node.Generate().Int64())
}

type HostSnowflakeMachineIDGenerator struct{}

func NewHostSnowflakeMachineIDGenerator() MachineIDGenerator {
	return &HostSnowflakeMachineIDGenerator{}
}

func (m *HostSnowflakeMachineIDGenerator) MustGenerateMachineID() uint16 {
	localIPs, err := getLocalIPs()
	if err != nil {
		panic(err)
	}
	if len(localIPs) == 0 {
		panic("empty ip address")
	}

	var ipBytes []byte
	for _, ip := range localIPs {
		ipBytes = append(ipBytes, []byte(ip)...)
	}
	return mustGetValidSnowflakeMachineID(ipBytes)
}

// mustGetValidSnowflakeMachineID get integer max to 10 bits
func mustGetValidSnowflakeMachineID(rawBytes []byte) uint16 {
	hash := sha256.New()
	if _, err := hash.Write(rawBytes); err != nil {
		panic(err)
	}

	machineIDBytes := make([]byte, 2)
	copy(machineIDBytes, hash.Sum(nil))
	machineIDBytes[0] = machineIDBytes[0] & 3 // only remain the last 2 bits
	// the whole machineIDBytes has only 10 last valid bits, parse the 10 bits to uint16
	return binary.BigEndian.Uint16(machineIDBytes)
}

func getLocalIPs() ([]string, error) {
	var ips []string
	interfaces, err := net.Interfaces()
	if err != nil {
		return nil, fmt.Errorf("get network interfaces error: %w", err)
	}

	for _, i := range interfaces {
		addrs, err := i.Addrs()
		if err != nil {
			return nil, fmt.Errorf("get addrs of network interface %s error: %w", i.Name, err)
		}
		for _, addr := range addrs {
			switch v := addr.(type) {
			case *net.IPNet:
				if !v.IP.IsLoopback() && v.IP.To4() != nil {
					ips = append(ips, v.IP.String())
				}
			case *net.IPAddr:
				if !v.IP.IsLoopback() && v.IP.To4() != nil {
					ips = append(ips, v.IP.String())
				}
			}
		}
	}
	return ips, nil
}