batch.go

package store

import (
	"context"
	"crypto/sha256"
	"fmt"

	ds "github.com/ipfs/go-datastore"

	"github.com/evstack/ev-node/types"
)

// DefaultBatch provides a batched write interface for atomic store operations
type DefaultBatch struct {
	store *DefaultStore
	batch ds.Batch
	ctx   context.Context
}

// NewBatch creates a new batch for atomic operations
func (s *DefaultStore) NewBatch(ctx context.Context) (Batch, error) {
	batch, err := s.db.Batch(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to create a new batch: %w", err)
	}

	return &DefaultBatch{
		store: s,
		batch: batch,
		ctx:   ctx,
	}, nil
}

// SetHeight sets the height in the batch
func (b *DefaultBatch) SetHeight(height uint64) error {
	currentHeight, err := b.store.Height(b.ctx)
	if err != nil {
		return err
	}
	if height <= currentHeight {
		return nil
	}

	heightBytes := encodeHeight(height)
	return b.batch.Put(b.ctx, ds.RawKey(getHeightKey()), heightBytes)
}

// SaveBlockData saves block data to the batch
func (b *DefaultBatch) SaveBlockData(header *types.SignedHeader, data *types.Data, signature *types.Signature) error {
	headerBlob, err := header.MarshalBinary()
	if err != nil {
		return fmt.Errorf("failed to marshal Header to binary: %w", err)
	}
	dataBlob, err := data.MarshalBinary()
	if err != nil {
		return fmt.Errorf("failed to marshal Data to binary: %w", err)
	}
	return b.SaveBlockDataFromBytes(header, headerBlob, dataBlob, signature)
}

// SaveBlockDataFromBytes saves pre-serialized block data to the batch.
// This avoids re-marshalling header and data when the caller already has the binary blobs.
func (b *DefaultBatch) SaveBlockDataFromBytes(header *types.SignedHeader, headerBlob, dataBlob []byte, signature *types.Signature) error {
	height := header.Height()
	signatureHash := *signature

	if err := b.batch.Put(b.ctx, ds.RawKey(getHeaderKey(height)), headerBlob); err != nil {
		return fmt.Errorf("failed to put header blob in batch: %w", err)
	}
	if err := b.batch.Put(b.ctx, ds.RawKey(getDataKey(height)), dataBlob); err != nil {
		return fmt.Errorf("failed to put data blob in batch: %w", err)
	}
	if err := b.batch.Put(b.ctx, ds.RawKey(getSignatureKey(height)), signatureHash[:]); err != nil {
		return fmt.Errorf("failed to put signature blob in batch: %w", err)
	}

	headerHash := sha256.Sum256(headerBlob)
	heightBytes := encodeHeight(height)
	if err := b.batch.Put(b.ctx, ds.RawKey(getIndexKey(headerHash[:])), heightBytes); err != nil {
		return fmt.Errorf("failed to put index key in batch: %w", err)
	}

	return nil
}

// UpdateState updates the state in the batch.
// Uses pooled State.MarshalBinary to reduce per-block allocations.
func (b *DefaultBatch) UpdateState(state types.State) error {
	height := state.LastBlockHeight

	data, err := state.MarshalBinary()
	if err != nil {
		return fmt.Errorf("failed to marshal state: %w", err)
	}

	return b.batch.Put(b.ctx, ds.RawKey(getStateAtHeightKey(height)), data)
}

// Commit commits all batched operations atomically
func (b *DefaultBatch) Commit() error {
	if err := b.batch.Commit(b.ctx); err != nil {
		return fmt.Errorf("failed to commit batch: %w", err)
	}
	return nil
}

// Delete adds a delete operation to the batch
func (b *DefaultBatch) Delete(key ds.Key) error {
	return b.batch.Delete(b.ctx, key)
}

// Put adds a put operation to the batch
func (b *DefaultBatch) Put(key ds.Key, value []byte) error {
	return b.batch.Put(b.ctx, key, value)
}