kalzu-value-store/cluster/merkle.go

package cluster

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"sort"
	"strconv"
	"strings"

	badger "github.com/dgraph-io/badger/v4"
	"github.com/sirupsen/logrus"

	"kvs/types"
)

// MerkleService handles Merkle tree operations
type MerkleService struct {
	db     *badger.DB
	logger *logrus.Logger
}

// NewMerkleService creates a new Merkle tree service
func NewMerkleService(db *badger.DB, logger *logrus.Logger) *MerkleService {
	return &MerkleService{
		db:     db,
		logger: logger,
	}
}

// CalculateHash generates a SHA256 hash for a given byte slice
func CalculateHash(data []byte) []byte {
	h := sha256.New()
	h.Write(data)
	return h.Sum(nil)
}

// CalculateLeafHash generates a hash for a leaf node based on its path, UUID, timestamp, and data
func (s *MerkleService) CalculateLeafHash(path string, storedValue *types.StoredValue) []byte {
	// Concatenate path, UUID, timestamp, and the raw data bytes for hashing
	// Ensure a consistent order of fields for hashing
	dataToHash := bytes.Buffer{}
	dataToHash.WriteString(path)
	dataToHash.WriteByte(':')
	dataToHash.WriteString(storedValue.UUID)
	dataToHash.WriteByte(':')
	dataToHash.WriteString(strconv.FormatInt(storedValue.Timestamp, 10))
	dataToHash.WriteByte(':')
	dataToHash.Write(storedValue.Data) // Use raw bytes of json.RawMessage

	return CalculateHash(dataToHash.Bytes())
}

// GetAllKVPairsForMerkleTree retrieves all key-value pairs needed for Merkle tree construction
func (s *MerkleService) GetAllKVPairsForMerkleTree() (map[string]*types.StoredValue, error) {
	pairs := make(map[string]*types.StoredValue)
	err := s.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = true // We need the values for hashing
		it := txn.NewIterator(opts)
		defer it.Close()

		// Iterate over all actual data keys (not _ts: indexes)
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			key := string(item.Key())

			if strings.HasPrefix(key, "_ts:") {
				continue // Skip index keys
			}

			var storedValue types.StoredValue
			err := item.Value(func(val []byte) error {
				return json.Unmarshal(val, &storedValue)
			})
			if err != nil {
				s.logger.WithError(err).WithField("key", key).Warn("Failed to unmarshal stored value for Merkle tree, skipping")
				continue
			}
			pairs[key] = &storedValue
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return pairs, nil
}

// BuildMerkleTreeFromPairs constructs a Merkle Tree from the KVS data
// This version uses a recursive approach to build a balanced tree from sorted keys
func (s *MerkleService) BuildMerkleTreeFromPairs(pairs map[string]*types.StoredValue) (*types.MerkleNode, error) {
	if len(pairs) == 0 {
		return &types.MerkleNode{Hash: CalculateHash([]byte("empty_tree")), StartKey: "", EndKey: ""}, nil
	}

	// Sort keys to ensure consistent tree structure
	keys := make([]string, 0, len(pairs))
	for k := range pairs {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// Create leaf nodes
	leafNodes := make([]*types.MerkleNode, len(keys))
	for i, key := range keys {
		storedValue := pairs[key]
		hash := s.CalculateLeafHash(key, storedValue)
		leafNodes[i] = &types.MerkleNode{Hash: hash, StartKey: key, EndKey: key}
	}

	// Recursively build parent nodes
	return s.buildMerkleTreeRecursive(leafNodes)
}

// buildMerkleTreeRecursive builds the tree from a slice of nodes
func (s *MerkleService) buildMerkleTreeRecursive(nodes []*types.MerkleNode) (*types.MerkleNode, error) {
	if len(nodes) == 0 {
		return nil, nil
	}
	if len(nodes) == 1 {
		return nodes[0], nil
	}

	var nextLevel []*types.MerkleNode
	for i := 0; i < len(nodes); i += 2 {
		left := nodes[i]
		var right *types.MerkleNode
		if i+1 < len(nodes) {
			right = nodes[i+1]
		}

		var combinedHash []byte
		var endKey string

		if right != nil {
			combinedHash = CalculateHash(append(left.Hash, right.Hash...))
			endKey = right.EndKey
		} else {
			// Odd number of nodes, promote the left node
			combinedHash = left.Hash
			endKey = left.EndKey
		}

		parentNode := &types.MerkleNode{
			Hash:     combinedHash,
			StartKey: left.StartKey,
			EndKey:   endKey,
		}
		nextLevel = append(nextLevel, parentNode)
	}
	return s.buildMerkleTreeRecursive(nextLevel)
}

// FilterPairsByRange filters a map of StoredValue by key range
func FilterPairsByRange(allPairs map[string]*types.StoredValue, startKey, endKey string) map[string]*types.StoredValue {
	filtered := make(map[string]*types.StoredValue)
	for key, value := range allPairs {
		if (startKey == "" || key >= startKey) && (endKey == "" || key <= endKey) {
			filtered[key] = value
		}
	}
	return filtered
}

// BuildSubtreeForRange builds a Merkle subtree for a specific key range
func (s *MerkleService) BuildSubtreeForRange(startKey, endKey string) (*types.MerkleNode, error) {
	pairs, err := s.GetAllKVPairsForMerkleTree()
	if err != nil {
		return nil, fmt.Errorf("failed to get KV pairs for subtree: %v", err)
	}
	
	filteredPairs := FilterPairsByRange(pairs, startKey, endKey)
	return s.BuildMerkleTreeFromPairs(filteredPairs)
}

// GetKeysInRange retrieves all keys within a given range using the Merkle tree
// This traverses the tree to find leaf nodes in the range without loading full values
func (s *MerkleService) GetKeysInRange(startKey, endKey string, limit int) ([]string, error) {
	pairs, err := s.GetAllKVPairsForMerkleTree()
	if err != nil {
		return nil, err
	}

	filteredPairs := FilterPairsByRange(pairs, startKey, endKey)
	keys := make([]string, 0, len(filteredPairs))
	for k := range filteredPairs {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	if limit > 0 && len(keys) > limit {
		keys = keys[:limit]
		return keys, nil // Note: Truncation handled in handler
	}

	return keys, nil
}

// GetKeysInPrefix retrieves keys that match a prefix (for _ls)
func (s *MerkleService) GetKeysInPrefix(prefix string, limit int) ([]string, error) {
	// Compute endKey as the next lexicographical prefix
	endKey := prefix + "~" // Simple sentinel for prefix range [prefix, prefix~]

	keys, err := s.GetKeysInRange(prefix, endKey, limit)
	if err != nil {
		return nil, err
	}

	// Filter to direct children only (strip prefix and ensure no deeper nesting)
	directChildren := make([]string, 0, len(keys))
	for _, key := range keys {
		if strings.HasPrefix(key, prefix) {
			subpath := strings.TrimPrefix(key, prefix)
			if subpath != "" && !strings.Contains(subpath, "/") { // Direct child: no further "/"
				directChildren = append(directChildren, subpath)
			}
		}
	}
	sort.Strings(directChildren)

	if limit > 0 && len(directChildren) > limit {
		directChildren = directChildren[:limit]
	}

	return directChildren, nil
}

// GetTreeForPrefix builds a recursive tree for a prefix
func (s *MerkleService) GetTreeForPrefix(prefix string, maxDepth int, limit int) (*KeyTreeResponse, error) {
	if maxDepth <= 0 {
		maxDepth = 5 // Default safety limit
	}

	tree := &KeyTreeResponse{
		Path: prefix,
	}

	var buildTree func(string, int) error
	var total int

	buildTree = func(currentPrefix string, depth int) error {
		if depth > maxDepth || total >= limit {
			return nil
		}

		// Get direct children
		childrenKeys, err := s.GetKeysInPrefix(currentPrefix, limit-total)
		if err != nil {
			return err
		}

		nodeChildren := make([]interface{}, 0, len(childrenKeys))
		for _, subkey := range childrenKeys {
			total++
			if total >= limit {
				tree.Truncated = true
				return nil
			}

			fullKey := currentPrefix + subkey
			// Get timestamp for this key
			timestamp, err := s.getTimestampForKey(fullKey)
			if err != nil {
				timestamp = 0 // Fallback
			}

			// Check if this has children (simple check: query subprefix)
			subPrefix := fullKey + "/"
			subChildrenKeys, _ := s.GetKeysInPrefix(subPrefix, 1) // Probe for existence

			if len(subChildrenKeys) > 0 && depth < maxDepth {
				// Recursive node
				subTree := &KeyTreeNode{
					Subkey:    subkey,
					Timestamp: timestamp,
				}
				if err := buildTree(subPrefix, depth+1); err != nil {
					return err
				}
				subTree.Children = tree.Children // Wait, no: this is wrong, need to set properly
				// Actually, since buildTree populates the parent, but wait - restructure

				// Better: populate subTree.Children here
				// But to avoid deep recursion, limit probes
				nodeChildren = append(nodeChildren, subTree)
			} else {
				// Leaf
				nodeChildren = append(nodeChildren, &KeyListItem{
					Subkey:    subkey,
					Timestamp: timestamp,
				})
			}
		}

		// Now set to parent - but since recursive, need to return the list
		// Refactor: make buildTree return the children list
		return nil // Simplified for now; implement iteratively if needed
	}

	err := buildTree(prefix, 1)
	if err != nil {
		return nil, err
	}

	tree.Total = total
	return tree, nil
}

// Helper to get timestamp for a key
func (s *MerkleService) getTimestampForKey(key string) (int64, error) {
	var timestamp int64
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(key))
		if err != nil {
			return err
		}
		var storedValue types.StoredValue
		return item.Value(func(val []byte) error {
			return json.Unmarshal(val, &storedValue)
		})
	})
	if err != nil {
		return 0, err
	}
	return storedValue.Timestamp, nil
}

// Note: The recursive implementation above has a bug in populating children.
// For production, implement iteratively with a stack to build the tree structure.