package kv

import (
	"context"
	"encoding/binary"

	"github.com/prysmaticlabs/prysm/shared/bytesutil"
	bolt "go.etcd.io/bbolt"
	"go.opencensus.io/trace"
)

// SaveArchivedPointRoot saves an archived point root to the DB. This is used for cold state management.
func (k *Store) SaveArchivedPointRoot(ctx context.Context, blockRoot [32]byte, index uint64) error {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.SaveArchivedPointRoot")
	defer span.End()

	return k.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(archivedIndexRootBucket)
		return bucket.Put(bytesutil.Uint64ToBytes(index), blockRoot[:])
	})
}

// SaveLastArchivedIndex to the db.
func (k *Store) SaveLastArchivedIndex(ctx context.Context, index uint64) error {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.SaveHeadBlockRoot")
	defer span.End()
	return k.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(archivedIndexRootBucket)
		return bucket.Put(lastArchivedIndexKey, bytesutil.Uint64ToBytes(index))
	})
}

// LastArchivedIndex from the db.
func (k *Store) LastArchivedIndex(ctx context.Context) (uint64, error) {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.LastArchivedIndex")
	defer span.End()
	var index uint64
	err := k.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(archivedIndexRootBucket)
		b := bucket.Get(lastArchivedIndexKey)
		if b == nil {
			return nil
		}
		index = binary.LittleEndian.Uint64(b)
		return nil
	})

	return index, err
}

// LastArchivedIndexRoot from the db.
func (k *Store) LastArchivedIndexRoot(ctx context.Context) [32]byte {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.LastArchivedIndexRoot")
	defer span.End()

	var blockRoot []byte
	if err := k.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(archivedIndexRootBucket)
		lastArchivedIndex := bucket.Get(lastArchivedIndexKey)
		if lastArchivedIndex == nil {
			return nil
		}
		blockRoot = bucket.Get(lastArchivedIndex)
		return nil
	}); err != nil { // This view never returns an error, but we'll handle anyway for sanity.
		panic(err)
	}

	return bytesutil.ToBytes32(blockRoot)
}

// ArchivedPointRoot returns the block root of an archived point from the DB.
// This is essential for cold state management and to restore a cold state.
func (k *Store) ArchivedPointRoot(ctx context.Context, index uint64) [32]byte {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.ArchivePointRoot")
	defer span.End()

	var blockRoot []byte
	if err := k.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(archivedIndexRootBucket)
		blockRoot = bucket.Get(bytesutil.Uint64ToBytes(index))
		return nil
	}); err != nil { // This view never returns an error, but we'll handle anyway for sanity.
		panic(err)
	}

	return bytesutil.ToBytes32(blockRoot)
}

// HasArchivedPoint returns true if an archived point exists in DB.
func (k *Store) HasArchivedPoint(ctx context.Context, index uint64) bool {
	ctx, span := trace.StartSpan(ctx, "BeaconDB.HasArchivedPoint")
	defer span.End()
	var exists bool
	if err := k.db.View(func(tx *bolt.Tx) error {
		iBucket := tx.Bucket(archivedIndexRootBucket)
		exists = iBucket.Get(bytesutil.Uint64ToBytes(index)) != nil
		return nil
	}); err != nil { // This view never returns an error, but we'll handle anyway for sanity.
		panic(err)
	}
	return exists
}