package blossomsub

import (
	"sync"

	pb "source.quilibrium.com/quilibrium/monorepo/go-libp2p-blossomsub/pb"
)

// msgIDGenerator handles computing IDs for msgs
// It allows setting custom generators(MsgIdFunction) per bitmask
type msgIDGenerator struct {
	Default MsgIdFunction

	bitmaskGensLk sync.RWMutex
	bitmaskGens   map[string]MsgIdFunction
}

func newMsgIdGenerator() *msgIDGenerator {
	return &msgIDGenerator{
		Default:     DefaultMsgIdFn,
		bitmaskGens: make(map[string]MsgIdFunction),
	}
}

// Set sets custom id generator(MsgIdFunction) for bitmask.
func (m *msgIDGenerator) Set(bitmask []byte, gen MsgIdFunction) {
	m.bitmaskGensLk.Lock()
	m.bitmaskGens[string(bitmask)] = gen
	m.bitmaskGensLk.Unlock()
}

// ID computes ID for the msg or short-circuits with the cached value.
func (m *msgIDGenerator) ID(msg *Message) []byte {
	if len(msg.ID) != 0 {
		return msg.ID
	}

	msg.ID = m.RawID(msg.Message)
	return msg.ID
}

// RawID computes ID for the proto 'msg'.
func (m *msgIDGenerator) RawID(msg *pb.Message) []byte {
	m.bitmaskGensLk.RLock()
	gen, ok := m.bitmaskGens[string(msg.GetBitmask())]
	m.bitmaskGensLk.RUnlock()
	if !ok {
		gen = m.Default
	}

	return gen(msg)
}