package store

import (
	"encoding/binary"

	"github.com/cockroachdb/pebble"
	"github.com/pkg/errors"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"
	"source.quilibrium.com/quilibrium/monorepo/node/protobufs"
)

type KeyStore interface {
	NewTransaction() (Transaction, error)
	StageProvingKey(provingKey *protobufs.ProvingKeyAnnouncement) error
	IncludeProvingKey(
		inclusionCommitment *protobufs.InclusionCommitment,
		txn Transaction,
	) error
	GetStagedProvingKey(
		provingKey []byte,
	) (*protobufs.ProvingKeyAnnouncement, error)
	GetProvingKey(provingKey []byte) (*protobufs.InclusionCommitment, error)
	GetKeyBundle(
		provingKey []byte,
		frameNumber uint64,
	) (*protobufs.InclusionCommitment, error)
	GetLatestKeyBundle(provingKey []byte) (*protobufs.InclusionCommitment, error)
	PutKeyBundle(
		provingKey []byte,
		keyBundleCommitment *protobufs.InclusionCommitment,
		txn Transaction,
	) error
	RangeProvingKeys() (*PebbleProvingKeyIterator, error)
	RangeStagedProvingKeys() (*PebbleStagedProvingKeyIterator, error)
	RangeKeyBundleKeys(provingKey []byte) (*PebbleKeyBundleIterator, error)
}

type PebbleKeyStore struct {
	db     *pebble.DB
	logger *zap.Logger
}

type PebbleProvingKeyIterator struct {
	i *pebble.Iterator
}

type PebbleStagedProvingKeyIterator struct {
	i *pebble.Iterator
}

type PebbleKeyBundleIterator struct {
	i *pebble.Iterator
}

var pki = (*PebbleProvingKeyIterator)(nil)
var spki = (*PebbleStagedProvingKeyIterator)(nil)
var kbi = (*PebbleKeyBundleIterator)(nil)
var _ Iterator[*protobufs.InclusionCommitment] = pki
var _ Iterator[*protobufs.ProvingKeyAnnouncement] = spki
var _ Iterator[*protobufs.InclusionCommitment] = kbi
var _ KeyStore = (*PebbleKeyStore)(nil)

func (p *PebbleProvingKeyIterator) First() bool {
	return p.i.First()
}

func (p *PebbleProvingKeyIterator) Next() bool {
	return p.i.Next()
}

func (p *PebbleProvingKeyIterator) Valid() bool {
	return p.i.Valid()
}

func (p *PebbleProvingKeyIterator) Value() (
	*protobufs.InclusionCommitment,
	error,
) {
	if !p.i.Valid() {
		return nil, ErrNotFound
	}

	value := p.i.Value()
	frame := &protobufs.InclusionCommitment{}
	if err := proto.Unmarshal(value, frame); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get proving key iterator value",
		)
	}

	return frame, nil
}

func (p *PebbleProvingKeyIterator) Close() error {
	return errors.Wrap(p.i.Close(), "closing iterator")
}

func (p *PebbleStagedProvingKeyIterator) First() bool {
	return p.i.First()
}

func (p *PebbleStagedProvingKeyIterator) Next() bool {
	return p.i.Next()
}

func (p *PebbleStagedProvingKeyIterator) Valid() bool {
	return p.i.Valid()
}

func (p *PebbleStagedProvingKeyIterator) Value() (
	*protobufs.ProvingKeyAnnouncement,
	error,
) {
	if !p.i.Valid() {
		return nil, ErrNotFound
	}

	value := p.i.Value()
	frame := &protobufs.ProvingKeyAnnouncement{}
	if err := proto.Unmarshal(value, frame); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get staged proving key iterator value",
		)
	}

	return frame, nil
}

func (p *PebbleStagedProvingKeyIterator) Close() error {
	return errors.Wrap(p.i.Close(), "closing iterator")
}

func (p *PebbleKeyBundleIterator) First() bool {
	return p.i.First()
}

func (p *PebbleKeyBundleIterator) Next() bool {
	return p.i.Next()
}

func (p *PebbleKeyBundleIterator) Valid() bool {
	return p.i.Valid()
}

func (p *PebbleKeyBundleIterator) Value() (
	*protobufs.InclusionCommitment,
	error,
) {
	if !p.i.Valid() {
		return nil, ErrNotFound
	}

	value := p.i.Value()
	frame := &protobufs.InclusionCommitment{}
	if err := proto.Unmarshal(value, frame); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get key bundle iterator value",
		)
	}

	return frame, nil
}

func (p *PebbleKeyBundleIterator) Close() error {
	return errors.Wrap(p.i.Close(), "closing iterator")
}

func NewPebbleKeyStore(db *pebble.DB, logger *zap.Logger) *PebbleKeyStore {
	return &PebbleKeyStore{
		db,
		logger,
	}
}

const (
	PROVING_KEY               = 0x01
	PROVING_KEY_STAGED        = 0x02
	KEY_BUNDLE                = 0x03
	KEY_DATA                  = 0x00
	KEY_BUNDLE_INDEX_EARLIEST = 0x10
	KEY_BUNDLE_INDEX_LATEST   = 0x20
)

func provingKeyKey(provingKey []byte) []byte {
	key := []byte{PROVING_KEY, KEY_DATA}
	key = append(key, provingKey...)
	return key
}

func stagedProvingKeyKey(provingKey []byte) []byte {
	key := []byte{PROVING_KEY_STAGED, KEY_DATA}
	key = append(key, provingKey...)
	return key
}

func keyBundleKey(provingKey []byte, frameNumber uint64) []byte {
	key := []byte{KEY_BUNDLE, KEY_DATA}
	key = append(key, provingKey...)
	key = binary.BigEndian.AppendUint64(key, frameNumber)
	return key
}

func keyBundleLatestKey(provingKey []byte) []byte {
	key := []byte{KEY_BUNDLE, KEY_BUNDLE_INDEX_LATEST}
	key = append(key, provingKey...)
	return key
}

func keyBundleEarliestKey(provingKey []byte) []byte {
	key := []byte{KEY_BUNDLE, KEY_BUNDLE_INDEX_EARLIEST}
	key = append(key, provingKey...)
	return key
}

func (p *PebbleKeyStore) NewTransaction() (Transaction, error) {
	return &PebbleTransaction{
		b: p.db.NewBatch(),
	}, nil
}

// Stages a proving key for later inclusion on proof of meaningful work.
// Does not verify, upstream callers must verify.
func (p *PebbleKeyStore) StageProvingKey(
	provingKey *protobufs.ProvingKeyAnnouncement,
) error {
	data, err := proto.Marshal(provingKey)
	if err != nil {
		return errors.Wrap(err, "stage proving key")
	}

	err = p.db.Set(
		stagedProvingKeyKey(provingKey.PublicKey()),
		data,
		&pebble.WriteOptions{
			Sync: true,
		},
	)
	if err != nil {
		return errors.Wrap(err, "stage proving key")
	}

	return nil
}

// Includes a proving key with an inclusion commitment. If a proving key is
// staged, promotes it by including it in the primary key store and deletes the
// staged key.
func (p *PebbleKeyStore) IncludeProvingKey(
	inclusionCommitment *protobufs.InclusionCommitment,
	txn Transaction,
) error {
	provingKey := &protobufs.ProvingKeyAnnouncement{}
	if err := proto.Unmarshal(inclusionCommitment.Data, provingKey); err != nil {
		return errors.Wrap(err, "include proving key")
	}

	if err := provingKey.Verify(); err != nil {
		return errors.Wrap(err, "include proving key")
	}

	data, err := proto.Marshal(inclusionCommitment)
	if err != nil {
		return errors.Wrap(err, "include proving key")
	}

	txn.Set(
		provingKeyKey(provingKey.PublicKey()),
		data,
	)

	staged, closer, err := p.db.Get(stagedProvingKeyKey(provingKey.PublicKey()))
	if err != nil && !errors.Is(err, ErrNotFound) {
		return errors.Wrap(err, "include proving key")
	}

	if staged != nil {
		if err := txn.Delete(
			stagedProvingKeyKey(provingKey.PublicKey()),
		); err != nil {
			return errors.Wrap(err, "include proving key")
		}
	}
	if err := closer.Close(); err != nil {
		return errors.Wrap(err, "include proving key")
	}

	return nil
}

func (p *PebbleKeyStore) GetStagedProvingKey(
	provingKey []byte,
) (*protobufs.ProvingKeyAnnouncement, error) {
	data, closer, err := p.db.Get(stagedProvingKeyKey(provingKey))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get staged proving key")
	}

	stagedKey := &protobufs.ProvingKeyAnnouncement{}
	if err = proto.Unmarshal(data, stagedKey); err != nil {
		return nil, errors.Wrap(err, "get staged proving key")
	}

	if err := closer.Close(); err != nil {
		return nil, errors.Wrap(err, "get staged proving key")
	}

	return stagedKey, nil
}

// Returns the latest key bundle for a given proving key.
func (p *PebbleKeyStore) GetLatestKeyBundle(
	provingKey []byte,
) (*protobufs.InclusionCommitment, error) {
	value, closer, err := p.db.Get(keyBundleLatestKey(provingKey))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get latest key bundle")
	}
	frameNumber := binary.BigEndian.Uint64(value)

	if err := closer.Close(); err != nil {
		return nil, errors.Wrap(err, "get latest key bundle")
	}

	value, closer, err = p.db.Get(keyBundleKey(provingKey, frameNumber))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get latest key bundle")
	}

	defer closer.Close()

	announcement := &protobufs.InclusionCommitment{}
	if err := proto.Unmarshal(value, announcement); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get latest key bundle",
		)
	}

	return announcement, nil
}

// Retrieves the specific key bundle included at a given frame number.
func (p *PebbleKeyStore) GetKeyBundle(
	provingKey []byte,
	frameNumber uint64,
) (*protobufs.InclusionCommitment, error) {
	value, closer, err := p.db.Get(keyBundleKey(provingKey, frameNumber))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get key bundle")
	}

	defer closer.Close()

	announcement := &protobufs.InclusionCommitment{}
	if err := proto.Unmarshal(value, announcement); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get key bundle",
		)
	}

	return announcement, nil
}

// Retrieves an included proving key, returns ErrNotFound if not present.
func (p *PebbleKeyStore) GetProvingKey(
	provingKey []byte,
) (*protobufs.InclusionCommitment, error) {
	value, closer, err := p.db.Get(provingKeyKey(provingKey))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get proving key")
	}

	defer closer.Close()

	announcement := &protobufs.InclusionCommitment{}
	if err := proto.Unmarshal(value, announcement); err != nil {
		return nil, errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"get proving key",
		)
	}

	return announcement, nil
}

// Inserts a key bundle with inclusion commitment. Does not verify, upstream
// callers must perform the verification.
func (p *PebbleKeyStore) PutKeyBundle(
	provingKey []byte,
	keyBundle *protobufs.InclusionCommitment,
	txn Transaction,
) error {
	data, err := proto.Marshal(keyBundle)
	if err != nil {
		return errors.Wrap(
			errors.Wrap(err, ErrInvalidData.Error()),
			"put key bundle",
		)
	}

	frameNumberBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(frameNumberBytes, keyBundle.FrameNumber)

	if err = txn.Set(
		keyBundleKey(provingKey, keyBundle.FrameNumber),
		data,
	); err != nil {
		return errors.Wrap(err, "put key bundle")
	}

	_, closer, err := p.db.Get(keyBundleEarliestKey(provingKey))
	if err != nil {
		if !errors.Is(err, pebble.ErrNotFound) {
			return errors.Wrap(err, "put key bundle")
		}

		if err = txn.Set(
			keyBundleEarliestKey(provingKey),
			frameNumberBytes,
		); err != nil {
			return errors.Wrap(err, "put key bundle")
		}
	}

	if err == nil && closer != nil {
		closer.Close()
	}

	if err = txn.Set(
		keyBundleLatestKey(provingKey),
		frameNumberBytes,
	); err != nil {
		return errors.Wrap(err, "put key bundle")
	}

	return nil
}

func (p *PebbleKeyStore) RangeProvingKeys() (*PebbleProvingKeyIterator, error) {
	iter := p.db.NewIter(&pebble.IterOptions{
		LowerBound: provingKeyKey([]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00,
		}),
		UpperBound: provingKeyKey([]byte{
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff,
		}),
	})

	return &PebbleProvingKeyIterator{i: iter}, nil
}

func (p *PebbleKeyStore) RangeStagedProvingKeys() (
	*PebbleStagedProvingKeyIterator,
	error,
) {
	iter := p.db.NewIter(&pebble.IterOptions{
		LowerBound: stagedProvingKeyKey([]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00,
		}),
		UpperBound: stagedProvingKeyKey([]byte{
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff,
		}),
	})

	return &PebbleStagedProvingKeyIterator{i: iter}, nil
}

func (p *PebbleKeyStore) RangeKeyBundleKeys(provingKey []byte) (
	*PebbleKeyBundleIterator,
	error,
) {
	iter := p.db.NewIter(&pebble.IterOptions{
		LowerBound: keyBundleKey(provingKey, 0),
		UpperBound: keyBundleKey(provingKey, 0xffffffffffffffff),
	})

	return &PebbleKeyBundleIterator{i: iter}, nil
}