package pstoremem

import (
	"errors"
	"sync"

	ic "github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/peer"
	pstore "github.com/libp2p/go-libp2p/core/peerstore"
)

type memoryKeyBook struct {
	sync.RWMutex // same lock. wont happen a ton.
	pks          map[peer.ID]ic.PubKey
	sks          map[peer.ID]ic.PrivKey
}

var _ pstore.KeyBook = (*memoryKeyBook)(nil)

func NewKeyBook() *memoryKeyBook {
	return &memoryKeyBook{
		pks: map[peer.ID]ic.PubKey{},
		sks: map[peer.ID]ic.PrivKey{},
	}
}

func (mkb *memoryKeyBook) PeersWithKeys() peer.IDSlice {
	mkb.RLock()
	ps := make(peer.IDSlice, 0, len(mkb.pks)+len(mkb.sks))
	for p := range mkb.pks {
		ps = append(ps, p)
	}
	for p := range mkb.sks {
		if _, found := mkb.pks[p]; !found {
			ps = append(ps, p)
		}
	}
	mkb.RUnlock()
	return ps
}

func (mkb *memoryKeyBook) PubKey(p peer.ID) ic.PubKey {
	mkb.RLock()
	pk := mkb.pks[p]
	mkb.RUnlock()
	if pk != nil {
		return pk
	}
	pk, err := p.ExtractPublicKey()
	if err == nil {
		mkb.Lock()
		mkb.pks[p] = pk
		mkb.Unlock()
	}
	return pk
}

func (mkb *memoryKeyBook) AddPubKey(p peer.ID, pk ic.PubKey) error {
	// check it's correct first
	if !p.MatchesPublicKey(pk) {
		return errors.New("ID does not match PublicKey")
	}

	mkb.Lock()
	mkb.pks[p] = pk
	mkb.Unlock()
	return nil
}

func (mkb *memoryKeyBook) PrivKey(p peer.ID) ic.PrivKey {
	mkb.RLock()
	defer mkb.RUnlock()
	return mkb.sks[p]
}

func (mkb *memoryKeyBook) AddPrivKey(p peer.ID, sk ic.PrivKey) error {
	if sk == nil {
		return errors.New("sk is nil (PrivKey)")
	}

	// check it's correct first
	if !p.MatchesPrivateKey(sk) {
		return errors.New("ID does not match PrivateKey")
	}

	mkb.Lock()
	mkb.sks[p] = sk
	mkb.Unlock()
	return nil
}

func (mkb *memoryKeyBook) RemovePeer(p peer.ID) {
	mkb.Lock()
	delete(mkb.sks, p)
	delete(mkb.pks, p)
	mkb.Unlock()
}