network/pkg/storage/service.go

package storage

import (
	"database/sql"
	"fmt"
	"io"
	"sync"

	"github.com/libp2p/go-libp2p/core/network"
	"go.uber.org/zap"
)

// Service provides distributed storage functionality using RQLite
type Service struct {
	logger *zap.Logger
	db     *sql.DB
	mu     sync.RWMutex
}

// NewService creates a new storage service backed by RQLite
func NewService(db *sql.DB, logger *zap.Logger) (*Service, error) {
	service := &Service{
		logger: logger,
		db:     db,
	}

	// Initialize storage tables
	if err := service.initTables(); err != nil {
		return nil, fmt.Errorf("failed to initialize storage tables: %w", err)
	}

	return service, nil
}

// initTables creates the necessary tables for key-value storage
func (s *Service) initTables() error {
	// Create storage table with namespace support
	createTableSQL := `
		CREATE TABLE IF NOT EXISTS kv_storage (
			namespace TEXT NOT NULL,
			key TEXT NOT NULL,
			value BLOB NOT NULL,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY (namespace, key)
		)
	`

	// Create index for faster queries
	createIndexSQL := `
		CREATE INDEX IF NOT EXISTS idx_kv_storage_namespace_key 
		ON kv_storage(namespace, key)
	`

	if _, err := s.db.Exec(createTableSQL); err != nil {
		return fmt.Errorf("failed to create storage table: %w", err)
	}

	if _, err := s.db.Exec(createIndexSQL); err != nil {
		return fmt.Errorf("failed to create storage index: %w", err)
	}

	s.logger.Info("Storage tables initialized successfully")
	return nil
}

// HandleStorageStream handles incoming storage protocol streams
func (s *Service) HandleStorageStream(stream network.Stream) {
	defer stream.Close()

	// Read request
	data, err := io.ReadAll(stream)
	if err != nil {
		s.logger.Error("Failed to read storage request", zap.Error(err))
		return
	}

	var request StorageRequest
	if err := request.Unmarshal(data); err != nil {
		s.logger.Error("Failed to unmarshal storage request", zap.Error(err))
		return
	}

	// Process request
	response := s.processRequest(&request)

	// Send response
	responseData, err := response.Marshal()
	if err != nil {
		s.logger.Error("Failed to marshal storage response", zap.Error(err))
		return
	}

	if _, err := stream.Write(responseData); err != nil {
		s.logger.Error("Failed to write storage response", zap.Error(err))
		return
	}

	s.logger.Debug("Handled storage request",
		zap.String("type", string(request.Type)),
		zap.String("key", request.Key),
		zap.String("namespace", request.Namespace),
		zap.Bool("success", response.Success),
	)
}

// processRequest processes a storage request and returns a response
func (s *Service) processRequest(req *StorageRequest) *StorageResponse {
	switch req.Type {
	case MessageTypePut:
		return s.handlePut(req)
	case MessageTypeGet:
		return s.handleGet(req)
	case MessageTypeDelete:
		return s.handleDelete(req)
	case MessageTypeList:
		return s.handleList(req)
	case MessageTypeExists:
		return s.handleExists(req)
	default:
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("unknown message type: %s", req.Type),
		}
	}
}

// handlePut stores a key-value pair
func (s *Service) handlePut(req *StorageRequest) *StorageResponse {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Use REPLACE to handle both insert and update
	query := `
		REPLACE INTO kv_storage (namespace, key, value, updated_at) 
		VALUES (?, ?, ?, CURRENT_TIMESTAMP)
	`

	_, err := s.db.Exec(query, req.Namespace, req.Key, req.Value)
	if err != nil {
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("failed to store key: %v", err),
		}
	}

	s.logger.Debug("Stored key", zap.String("key", req.Key), zap.String("namespace", req.Namespace))
	return &StorageResponse{Success: true}
}

// handleGet retrieves a value by key
func (s *Service) handleGet(req *StorageRequest) *StorageResponse {
	s.mu.RLock()
	defer s.mu.RUnlock()

	query := `SELECT value FROM kv_storage WHERE namespace = ? AND key = ?`

	var value []byte
	err := s.db.QueryRow(query, req.Namespace, req.Key).Scan(&value)
	if err != nil {
		if err == sql.ErrNoRows {
			return &StorageResponse{
				Success: false,
				Error:   fmt.Sprintf("key not found: %s", req.Key),
			}
		}
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("failed to get key: %v", err),
		}
	}

	return &StorageResponse{
		Success: true,
		Value:   value,
	}
}

// handleDelete removes a key
func (s *Service) handleDelete(req *StorageRequest) *StorageResponse {
	s.mu.Lock()
	defer s.mu.Unlock()

	query := `DELETE FROM kv_storage WHERE namespace = ? AND key = ?`

	result, err := s.db.Exec(query, req.Namespace, req.Key)
	if err != nil {
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("failed to delete key: %v", err),
		}
	}

	rowsAffected, _ := result.RowsAffected()
	if rowsAffected == 0 {
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("key not found: %s", req.Key),
		}
	}

	s.logger.Debug("Deleted key", zap.String("key", req.Key), zap.String("namespace", req.Namespace))
	return &StorageResponse{Success: true}
}

// handleList lists keys with a prefix
func (s *Service) handleList(req *StorageRequest) *StorageResponse {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var query string
	var args []interface{}

	if req.Prefix == "" {
		// List all keys in namespace
		query = `SELECT key FROM kv_storage WHERE namespace = ?`
		args = []interface{}{req.Namespace}
	} else {
		// List keys with prefix
		query = `SELECT key FROM kv_storage WHERE namespace = ? AND key LIKE ?`
		args = []interface{}{req.Namespace, req.Prefix + "%"}
	}

	if req.Limit > 0 {
		query += ` LIMIT ?`
		args = append(args, req.Limit)
	}

	rows, err := s.db.Query(query, args...)
	if err != nil {
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("failed to query keys: %v", err),
		}
	}
	defer rows.Close()

	var keys []string
	for rows.Next() {
		var key string
		if err := rows.Scan(&key); err != nil {
			continue
		}
		keys = append(keys, key)
	}

	return &StorageResponse{
		Success: true,
		Keys:    keys,
	}
}

// handleExists checks if a key exists
func (s *Service) handleExists(req *StorageRequest) *StorageResponse {
	s.mu.RLock()
	defer s.mu.RUnlock()

	query := `SELECT 1 FROM kv_storage WHERE namespace = ? AND key = ? LIMIT 1`

	var exists int
	err := s.db.QueryRow(query, req.Namespace, req.Key).Scan(&exists)
	if err != nil {
		if err == sql.ErrNoRows {
			return &StorageResponse{
				Success: true,
				Exists:  false,
			}
		}
		return &StorageResponse{
			Success: false,
			Error:   fmt.Sprintf("failed to check key existence: %v", err),
		}
	}

	return &StorageResponse{
		Success: true,
		Exists:  true,
	}
}

// Close closes the storage service
func (s *Service) Close() error {
	// The database connection is managed elsewhere
	s.logger.Info("Storage service closed")
	return nil
}