bzzz/pkg/ucxi/storage.go

package ucxi

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// BasicContentStorage provides a basic file-system based implementation of ContentStorage
type BasicContentStorage struct {
	basePath string
	mutex    sync.RWMutex
}

// NewBasicContentStorage creates a new basic content storage
func NewBasicContentStorage(basePath string) (*BasicContentStorage, error) {
	// Ensure base directory exists
	if err := os.MkdirAll(basePath, 0755); err != nil {
		return nil, fmt.Errorf("failed to create storage directory: %w", err)
	}

	return &BasicContentStorage{
		basePath: basePath,
	}, nil
}

// Store stores content with the given key
func (s *BasicContentStorage) Store(ctx context.Context, key string, content *Content) error {
	if key == "" {
		return fmt.Errorf("key cannot be empty")
	}
	if content == nil {
		return fmt.Errorf("content cannot be nil")
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	// Generate file path
	filePath := s.getFilePath(key)

	// Ensure directory exists
	dir := filepath.Dir(filePath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("failed to create directory %s: %w", dir, err)
	}

	// Calculate checksum if not provided
	if content.Checksum == "" {
		hash := sha256.Sum256(content.Data)
		content.Checksum = hex.EncodeToString(hash[:])
	}

	// Serialize content to JSON
	data, err := json.MarshalIndent(content, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to serialize content: %w", err)
	}

	// Write to file
	if err := ioutil.WriteFile(filePath, data, 0644); err != nil {
		return fmt.Errorf("failed to write content file: %w", err)
	}

	return nil
}

// Retrieve retrieves content by key
func (s *BasicContentStorage) Retrieve(ctx context.Context, key string) (*Content, error) {
	if key == "" {
		return nil, fmt.Errorf("key cannot be empty")
	}

	s.mutex.RLock()
	defer s.mutex.RUnlock()

	filePath := s.getFilePath(key)

	// Check if file exists
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		return nil, fmt.Errorf("content not found for key: %s", key)
	}

	// Read file
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read content file: %w", err)
	}

	// Deserialize content
	var content Content
	if err := json.Unmarshal(data, &content); err != nil {
		return nil, fmt.Errorf("failed to deserialize content: %w", err)
	}

	// Verify checksum if available
	if content.Checksum != "" {
		hash := sha256.Sum256(content.Data)
		expectedChecksum := hex.EncodeToString(hash[:])
		if content.Checksum != expectedChecksum {
			return nil, fmt.Errorf("content checksum mismatch")
		}
	}

	return &content, nil
}

// Delete deletes content by key
func (s *BasicContentStorage) Delete(ctx context.Context, key string) error {
	if key == "" {
		return fmt.Errorf("key cannot be empty")
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	filePath := s.getFilePath(key)

	// Check if file exists
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		return fmt.Errorf("content not found for key: %s", key)
	}

	// Remove file
	if err := os.Remove(filePath); err != nil {
		return fmt.Errorf("failed to delete content file: %w", err)
	}

	// Try to remove empty directories
	s.cleanupEmptyDirs(filepath.Dir(filePath))

	return nil
}

// List lists all keys with the given prefix
func (s *BasicContentStorage) List(ctx context.Context, prefix string) ([]string, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	var keys []string

	// Walk through storage directory
	err := filepath.Walk(s.basePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Skip non-JSON files
		if !strings.HasSuffix(path, ".json") {
			return nil
		}

		// Convert file path back to key
		relPath, err := filepath.Rel(s.basePath, path)
		if err != nil {
			return err
		}

		// Remove .json extension
		key := strings.TrimSuffix(relPath, ".json")
		
		// Convert file path separators back to key format
		key = strings.ReplaceAll(key, string(filepath.Separator), "/")

		// Check prefix match
		if prefix == "" || strings.HasPrefix(key, prefix) {
			keys = append(keys, key)
		}

		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to list storage contents: %w", err)
	}

	return keys, nil
}

// getFilePath converts a storage key to a file path
func (s *BasicContentStorage) getFilePath(key string) string {
	// Sanitize key by replacing potentially problematic characters
	sanitized := strings.ReplaceAll(key, ":", "_")
	sanitized = strings.ReplaceAll(sanitized, "@", "_at_")
	sanitized = strings.ReplaceAll(sanitized, "/", string(filepath.Separator))
	
	return filepath.Join(s.basePath, sanitized+".json")
}

// cleanupEmptyDirs removes empty directories up the tree
func (s *BasicContentStorage) cleanupEmptyDirs(dir string) {
	// Don't remove the base directory
	if dir == s.basePath {
		return
	}

	// Try to remove directory if empty
	if err := os.Remove(dir); err == nil {
		// Successfully removed, try parent
		s.cleanupEmptyDirs(filepath.Dir(dir))
	}
}

// GetStorageStats returns statistics about the storage
func (s *BasicContentStorage) GetStorageStats() (map[string]interface{}, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	var fileCount int
	var totalSize int64

	err := filepath.Walk(s.basePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() && strings.HasSuffix(path, ".json") {
			fileCount++
			totalSize += info.Size()
		}

		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to calculate storage stats: %w", err)
	}

	return map[string]interface{}{
		"file_count": fileCount,
		"total_size": totalSize,
		"base_path":  s.basePath,
	}, nil
}

// Exists checks if content exists for the given key
func (s *BasicContentStorage) Exists(ctx context.Context, key string) (bool, error) {
	if key == "" {
		return false, fmt.Errorf("key cannot be empty")
	}

	filePath := s.getFilePath(key)
	
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	_, err := os.Stat(filePath)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("failed to check file existence: %w", err)
	}

	return true, nil
}

// Clear removes all content from storage
func (s *BasicContentStorage) Clear(ctx context.Context) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// Remove all contents of base directory
	entries, err := ioutil.ReadDir(s.basePath)
	if err != nil {
		return fmt.Errorf("failed to read storage directory: %w", err)
	}

	for _, entry := range entries {
		path := filepath.Join(s.basePath, entry.Name())
		if err := os.RemoveAll(path); err != nil {
			return fmt.Errorf("failed to remove %s: %w", path, err)
		}
	}

	return nil
}