bzzz/pkg/slurp/storage/batch_operations.go

package storage

import (
	"context"
	"fmt"
	"sync"
	"time"

	"chorus.services/bzzz/pkg/ucxl"
	slurpContext "chorus.services/bzzz/pkg/slurp/context"
)

// BatchOperationsImpl provides efficient batch operations for context storage
type BatchOperationsImpl struct {
	contextStore  *ContextStoreImpl
	batchSize     int
	maxConcurrency int
	operationTimeout time.Duration
}

// NewBatchOperations creates a new batch operations handler
func NewBatchOperations(contextStore *ContextStoreImpl, batchSize, maxConcurrency int, timeout time.Duration) *BatchOperationsImpl {
	return &BatchOperationsImpl{
		contextStore:     contextStore,
		batchSize:       batchSize,
		maxConcurrency:  maxConcurrency,
		operationTimeout: timeout,
	}
}

// BatchStore stores multiple contexts efficiently
func (cs *ContextStoreImpl) BatchStore(
	ctx context.Context,
	batch *BatchStoreRequest,
) (*BatchStoreResult, error) {
	start := time.Now()
	result := &BatchStoreResult{
		Errors:      make(map[string]error),
		ProcessedAt: time.Now(),
	}

	// Validate batch request
	if batch == nil || len(batch.Contexts) == 0 {
		return result, fmt.Errorf("empty batch request")
	}

	// Create worker pool for concurrent processing
	workerCount := cs.options.MaxConcurrentOps
	if len(batch.Contexts) < workerCount {
		workerCount = len(batch.Contexts)
	}

	// Channels for work distribution
	workCh := make(chan *BatchStoreWork, len(batch.Contexts))
	resultsCh := make(chan *BatchStoreWorkResult, len(batch.Contexts))

	// Start workers
	var wg sync.WaitGroup
	for i := 0; i < workerCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			cs.batchStoreWorker(ctx, workCh, resultsCh, batch)
		}()
	}

	// Send work to workers
	go func() {
		defer close(workCh)
		for i, contextItem := range batch.Contexts {
			work := &BatchStoreWork{
				Index:   i,
				Item:    contextItem,
				Timeout: cs.options.OperationTimeout,
			}
			workCh <- work
		}
	}()

	// Collect results
	go func() {
		wg.Wait()
		close(resultsCh)
	}()

	// Process results
	for workResult := range resultsCh {
		if workResult.Error != nil {
			result.ErrorCount++
			key := workResult.Item.Context.UCXLAddress.String()
			result.Errors[key] = workResult.Error
			
			if batch.FailOnError {
				// Cancel remaining operations
				result.ProcessingTime = time.Since(start)
				return result, fmt.Errorf("batch operation failed on context %s: %w", key, workResult.Error)
			}
		} else {
			result.SuccessCount++
		}
	}

	result.ProcessingTime = time.Since(start)
	return result, nil
}

// BatchRetrieve retrieves multiple contexts efficiently
func (cs *ContextStoreImpl) BatchRetrieve(
	ctx context.Context,
	batch *BatchRetrieveRequest,
) (*BatchRetrieveResult, error) {
	start := time.Now()
	result := &BatchRetrieveResult{
		Contexts:    make(map[string]*slurpContext.ContextNode),
		Errors:      make(map[string]error),
		ProcessedAt: time.Now(),
	}

	// Validate batch request
	if batch == nil || len(batch.Addresses) == 0 {
		return result, fmt.Errorf("empty batch request")
	}

	// Create worker pool for concurrent processing
	workerCount := cs.options.MaxConcurrentOps
	if len(batch.Addresses) < workerCount {
		workerCount = len(batch.Addresses)
	}

	// Channels for work distribution
	workCh := make(chan *BatchRetrieveWork, len(batch.Addresses))
	resultsCh := make(chan *BatchRetrieveWorkResult, len(batch.Addresses))

	// Start workers
	var wg sync.WaitGroup
	for i := 0; i < workerCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			cs.batchRetrieveWorker(ctx, workCh, resultsCh, batch)
		}()
	}

	// Send work to workers
	go func() {
		defer close(workCh)
		for i, address := range batch.Addresses {
			work := &BatchRetrieveWork{
				Index:   i,
				Address: address,
				Role:    batch.Role,
				Timeout: cs.options.OperationTimeout,
			}
			workCh <- work
		}
	}()

	// Collect results
	go func() {
		wg.Wait()
		close(resultsCh)
	}()

	// Process results
	for workResult := range resultsCh {
		addressStr := workResult.Address.String()
		
		if workResult.Error != nil {
			result.ErrorCount++
			result.Errors[addressStr] = workResult.Error
			
			if batch.FailOnError {
				// Cancel remaining operations
				result.ProcessingTime = time.Since(start)
				return result, fmt.Errorf("batch operation failed on address %s: %w", addressStr, workResult.Error)
			}
		} else {
			result.SuccessCount++
			result.Contexts[addressStr] = workResult.Context
		}
	}

	result.ProcessingTime = time.Since(start)
	return result, nil
}

// Worker functions and supporting types

type BatchStoreWork struct {
	Index   int
	Item    *ContextStoreItem
	Timeout time.Duration
}

type BatchStoreWorkResult struct {
	Index int
	Item  *ContextStoreItem
	Error error
}

type BatchRetrieveWork struct {
	Index   int
	Address ucxl.Address
	Role    string
	Timeout time.Duration
}

type BatchRetrieveWorkResult struct {
	Index   int
	Address ucxl.Address
	Context *slurpContext.ContextNode
	Error   error
}

// batchStoreWorker processes batch store work items
func (cs *ContextStoreImpl) batchStoreWorker(
	ctx context.Context,
	workCh <-chan *BatchStoreWork,
	resultsCh chan<- *BatchStoreWorkResult,
	batch *BatchStoreRequest,
) {
	for work := range workCh {
		result := &BatchStoreWorkResult{
			Index: work.Index,
			Item:  work.Item,
		}

		// Create timeout context for this operation
		workCtx, cancel := context.WithTimeout(ctx, work.Timeout)
		defer cancel()

		// Determine roles to use
		roles := work.Item.Roles
		if len(roles) == 0 {
			roles = batch.Roles // Use batch default roles
		}

		// Perform the store operation
		if batch.Transaction {
			// Use transaction if requested
			result.Error = cs.storeContextWithTransaction(workCtx, work.Item.Context, roles)
		} else {
			// Regular store
			result.Error = cs.StoreContext(workCtx, work.Item.Context, roles)
		}

		resultsCh <- result
	}
}

// batchRetrieveWorker processes batch retrieve work items
func (cs *ContextStoreImpl) batchRetrieveWorker(
	ctx context.Context,
	workCh <-chan *BatchRetrieveWork,
	resultsCh chan<- *BatchRetrieveWorkResult,
	batch *BatchRetrieveRequest,
) {
	for work := range workCh {
		result := &BatchRetrieveWorkResult{
			Index:   work.Index,
			Address: work.Address,
		}

		// Create timeout context for this operation
		workCtx, cancel := context.WithTimeout(ctx, work.Timeout)
		defer cancel()

		// Perform the retrieve operation
		contextNode, err := cs.RetrieveContext(workCtx, work.Address, work.Role)
		result.Context = contextNode
		result.Error = err

		resultsCh <- result
	}
}

// storeContextWithTransaction performs a store operation within a transaction context
func (cs *ContextStoreImpl) storeContextWithTransaction(
	ctx context.Context,
	node *slurpContext.ContextNode,
	roles []string,
) error {
	// This would integrate with a transaction manager if available
	// For now, it's the same as regular store but could be enhanced
	// with rollback capabilities
	return cs.StoreContext(ctx, node, roles)
}

// ListContexts lists contexts matching criteria with optimized querying
func (cs *ContextStoreImpl) ListContexts(
	ctx context.Context,
	criteria *ListCriteria,
) ([]*slurpContext.ContextNode, error) {
	start := time.Now()
	defer func() {
		cs.recordLatency("list", time.Since(start))
	}()

	// Use search index if available and appropriate
	if cs.options.IndexingEnabled && cs.shouldUseSearchIndex(criteria) {
		return cs.listContextsViaSearch(ctx, criteria)
	}

	// Fallback to storage enumeration
	return cs.listContextsViaStorage(ctx, criteria)
}

// SearchContexts searches contexts using query criteria with advanced features
func (cs *ContextStoreImpl) SearchContexts(
	ctx context.Context,
	query *SearchQuery,
) (*SearchResults, error) {
	start := time.Now()
	defer func() {
		cs.recordLatency("search", time.Since(start))
	}()

	// Validate search query
	if query == nil {
		return nil, fmt.Errorf("search query cannot be nil")
	}

	// Use primary search index
	indexName := "primary"
	if cs.indexManager != nil {
		indexes, err := cs.indexManager.ListIndexes(ctx)
		if err == nil && len(indexes) > 0 {
			indexName = indexes[0] // Use first available index
		}
	}

	// Perform search
	results, err := cs.indexManager.Search(ctx, indexName, query)
	if err != nil {
		return nil, fmt.Errorf("search failed: %w", err)
	}

	// Post-process results for role-based filtering
	if len(query.Roles) > 0 {
		results = cs.filterResultsByRole(ctx, results, query.Roles)
	}

	// Apply additional filters that couldn't be done at index level
	results = cs.applyPostSearchFilters(ctx, results, query)

	cs.recordOperation("search")
	return results, nil
}

// Helper methods for optimized listing and searching

func (cs *ContextStoreImpl) shouldUseSearchIndex(criteria *ListCriteria) bool {
	// Use search index if we have complex criteria that would benefit from indexing
	return len(criteria.Tags) > 0 ||
		len(criteria.Technologies) > 0 ||
		criteria.PathPattern != "" ||
		criteria.MinConfidence > 0
}

func (cs *ContextStoreImpl) listContextsViaSearch(
	ctx context.Context,
	criteria *ListCriteria,
) ([]*slurpContext.ContextNode, error) {
	// Convert list criteria to search query
	query := &SearchQuery{
		Tags:          criteria.Tags,
		Technologies:  criteria.Technologies,
		Roles:         criteria.Roles,
		MinConfidence: criteria.MinConfidence,
		Limit:         criteria.Limit,
		Offset:        criteria.Offset,
		SortBy:        criteria.SortBy,
		SortOrder:     criteria.SortOrder,
		IncludeStale:  criteria.IncludeStale,
	}

	// Add path pattern as scope if provided
	if criteria.PathPattern != "" {
		query.Scope = []string{criteria.PathPattern}
	}

	// Perform search
	searchResults, err := cs.SearchContexts(ctx, query)
	if err != nil {
		return nil, err
	}

	// Extract contexts from search results
	contexts := make([]*slurpContext.ContextNode, len(searchResults.Results))
	for i, result := range searchResults.Results {
		contexts[i] = result.Context
	}

	return contexts, nil
}

func (cs *ContextStoreImpl) listContextsViaStorage(
	ctx context.Context,
	criteria *ListCriteria,
) ([]*slurpContext.ContextNode, error) {
	// This would enumerate storage and apply filters
	// This is a simplified implementation - in practice, this would be more sophisticated
	var contexts []*slurpContext.ContextNode

	// For now, return empty list as this would require storage enumeration
	// In a real implementation, this would iterate through storage keys
	// and load contexts that match the criteria

	return contexts, nil
}

func (cs *ContextStoreImpl) filterResultsByRole(
	ctx context.Context,
	results *SearchResults,
	roles []string,
) *SearchResults {
	// Filter search results based on role access
	// This ensures users only see contexts they have access to
	filteredResults := make([]*SearchResult, 0, len(results.Results))

	for _, result := range results.Results {
		// Check if any of the requested roles can access this context
		hasAccess := false
		for _, role := range roles {
			if cs.options.EncryptionEnabled {
				storageKey := cs.generateStorageKey(result.Context.UCXLAddress)
				if canAccess, err := cs.encryptedStorage.CanAccess(ctx, storageKey, role); err == nil && canAccess {
					hasAccess = true
					break
				}
			} else {
				// For unencrypted storage, assume access is allowed
				hasAccess = true
				break
			}
		}

		if hasAccess {
			filteredResults = append(filteredResults, result)
		}
	}

	// Update result metadata
	results.Results = filteredResults
	results.TotalResults = int64(len(filteredResults))

	return results
}

func (cs *ContextStoreImpl) applyPostSearchFilters(
	ctx context.Context,
	results *SearchResults,
	query *SearchQuery,
) *SearchResults {
	// Apply filters that couldn't be applied at the search index level
	filteredResults := make([]*SearchResult, 0, len(results.Results))

	for _, result := range results.Results {
		include := true

		// Age filter
		if query.MaxAge != nil {
			age := time.Since(result.Context.GeneratedAt)
			if age > *query.MaxAge {
				include = false
			}
		}

		// File type filter (based on path extension)
		if len(query.FileTypes) > 0 {
			matchesFileType := false
			for _, fileType := range query.FileTypes {
				if strings.HasSuffix(result.Context.Path, "."+fileType) {
					matchesFileType = true
					break
				}
			}
			if !matchesFileType {
				include = false
			}
		}

		// Scope filter
		if len(query.Scope) > 0 {
			matchesScope := false
			for _, scope := range query.Scope {
				if strings.HasPrefix(result.Context.Path, scope) {
					matchesScope = true
					break
				}
			}
			if !matchesScope {
				include = false
			}
		}

		// Exclude scope filter
		if len(query.ExcludeScope) > 0 {
			for _, excludeScope := range query.ExcludeScope {
				if strings.HasPrefix(result.Context.Path, excludeScope) {
					include = false
					break
				}
			}
		}

		if include {
			filteredResults = append(filteredResults, result)
		}
	}

	// Update result metadata
	results.Results = filteredResults
	results.TotalResults = int64(len(filteredResults))

	return results
}