...

pkg/heroservices/billing/models/db.go

@@ -0,0 +1,546 @@
+package models
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/freeflowuniverse/heroagent/pkg/data/ourdb"
+	"github.com/freeflowuniverse/heroagent/pkg/data/radixtree"
+	"github.com/freeflowuniverse/heroagent/pkg/tools"
+)
+
+// DBStore represents the central database store for all models
+type DBStore struct {
+	DB   *ourdb.OurDB
+	Meta *radixtree.RadixTree
+}
+
+// NewDBStore creates a new database store
+func NewDBStore() (*DBStore, error) {
+	// Get home directory
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get home directory: %w", err)
+	}
+
+	// Create paths
+	dbPath := filepath.Join(homeDir, "hero", "var", "server", "users", "db")
+	metaPath := filepath.Join(homeDir, "hero", "var", "server", "users", "meta")
+
+	// Create directories if they don't exist
+	if err := os.MkdirAll(dbPath, 0755); err != nil {
+		return nil, fmt.Errorf("failed to create db directory: %w", err)
+	}
+	if err := os.MkdirAll(metaPath, 0755); err != nil {
+		return nil, fmt.Errorf("failed to create meta directory: %w", err)
+	}
+
+	// Create ourdb
+	dbConfig := ourdb.DefaultConfig()
+	dbConfig.Path = dbPath
+	dbConfig.IncrementalMode = true
+	db, err := ourdb.New(dbConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create ourdb: %w", err)
+	}
+
+	// Create radixtree
+	meta, err := radixtree.New(radixtree.NewArgs{
+		Path: metaPath,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create radixtree: %w", err)
+	}
+
+	return &DBStore{
+		DB:   db,
+		Meta: meta,
+	}, nil
+}
+
+// Close closes the database connections
+func (s *DBStore) Close() error {
+	err1 := s.DB.Close()
+	err2 := s.Meta.Close()
+
+	if err1 != nil {
+		return err1
+	}
+	return err2
+}
+
+// UserStore handles the storage and retrieval of User objects
+type UserStore struct {
+	Store *DBStore
+}
+
+// NewUserStore creates a new UserStore
+func NewUserStore() (*UserStore, *DBStore, error) {
+	store, err := NewDBStore()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return &UserStore{
+		Store: store,
+	}, store, nil
+}
+
+// Save saves a User to the database
+func (s *UserStore) Save(user *User) error {
+	// If ID is 0, this is a new user
+	if user.ID == 0 {
+		// Fix the key
+		fixedKey := tools.NameFix(user.Key)
+		if fixedKey == "" {
+			return errors.New("key cannot be empty")
+		}
+		user.Key = fixedKey
+
+		// Check if key already exists
+		existingID, err := s.GetIDByKey(user.Key)
+		if err == nil && existingID != 0 {
+			return fmt.Errorf("user with key %s already exists", user.Key)
+		}
+
+		// Get next ID
+		nextID, err := s.Store.DB.GetNextID()
+		if err != nil {
+			return fmt.Errorf("failed to get next ID: %w", err)
+		}
+		user.ID = nextID
+
+		// Save to ourdb
+		data := user.Serialize()
+		_, err = s.Store.DB.Set(ourdb.OurDBSetArgs{
+			Data: data,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to save user to ourdb: %w", err)
+		}
+
+		// Save key to radixtree
+		idBytes := make([]byte, 4)
+		binary.LittleEndian.PutUint32(idBytes, user.ID)
+		err = s.Store.Meta.Set(user.Key, idBytes)
+		if err != nil {
+			return fmt.Errorf("failed to save user key to radixtree: %w", err)
+		}
+	} else {
+		// Update existing user
+		// Get existing user to verify key
+		existingUser, err := s.GetByID(user.ID)
+		if err != nil {
+			return fmt.Errorf("failed to get existing user: %w", err)
+		}
+
+		// If key changed, update radixtree
+		if existingUser.Key != user.Key {
+			// Fix the new key
+			fixedKey := tools.NameFix(user.Key)
+			if fixedKey == "" {
+				return errors.New("key cannot be empty")
+			}
+			user.Key = fixedKey
+
+			// Check if new key already exists
+			existingID, err := s.GetIDByKey(user.Key)
+			if err == nil && existingID != 0 && existingID != user.ID {
+				return fmt.Errorf("user with key %s already exists", user.Key)
+			}
+
+			// Delete old key from radixtree
+			err = s.Store.Meta.Delete(existingUser.Key)
+			if err != nil {
+				return fmt.Errorf("failed to delete old user key from radixtree: %w", err)
+			}
+
+			// Save new key to radixtree
+			idBytes := make([]byte, 4)
+			binary.LittleEndian.PutUint32(idBytes, user.ID)
+			err = s.Store.Meta.Set(user.Key, idBytes)
+			if err != nil {
+				return fmt.Errorf("failed to save new user key to radixtree: %w", err)
+			}
+		}
+
+		// Save to ourdb
+		data := user.Serialize()
+		id := user.ID
+		_, err = s.Store.DB.Set(ourdb.OurDBSetArgs{
+			ID:   &id,
+			Data: data,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to update user in ourdb: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// GetByID retrieves a User by ID
+func (s *UserStore) GetByID(id uint32) (*User, error) {
+	data, err := s.Store.DB.Get(id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get user from ourdb: %w", err)
+	}
+
+	user, err := DeserializeUser(data)
+	if err != nil {
+		return nil, fmt.Errorf("failed to deserialize user: %w", err)
+	}
+
+	return user, nil
+}
+
+// GetByKey retrieves a User by key
+func (s *UserStore) GetByKey(key string) (*User, error) {
+	// Fix the key
+	fixedKey := tools.NameFix(key)
+	if fixedKey == "" {
+		return nil, errors.New("key cannot be empty")
+	}
+
+	// Get ID from radixtree
+	id, err := s.GetIDByKey(fixedKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get user ID from radixtree: %w", err)
+	}
+
+	// Get user from ourdb
+	return s.GetByID(id)
+}
+
+// GetIDByKey retrieves a user ID by key
+func (s *UserStore) GetIDByKey(key string) (uint32, error) {
+	// Get ID from radixtree
+	idBytes, err := s.Store.Meta.Get(key)
+	if err != nil {
+		return 0, fmt.Errorf("failed to get user ID from radixtree: %w", err)
+	}
+
+	if len(idBytes) != 4 {
+		return 0, fmt.Errorf("invalid ID bytes length: %d", len(idBytes))
+	}
+
+	return binary.LittleEndian.Uint32(idBytes), nil
+}
+
+// Delete deletes a User by ID
+func (s *UserStore) Delete(id uint32) error {
+	// Get user to get key
+	user, err := s.GetByID(id)
+	if err != nil {
+		return fmt.Errorf("failed to get user: %w", err)
+	}
+
+	// Delete from radixtree
+	err = s.Store.Meta.Delete(user.Key)
+	if err != nil {
+		return fmt.Errorf("failed to delete user key from radixtree: %w", err)
+	}
+
+	// Delete from ourdb
+	err = s.Store.DB.Delete(id)
+	if err != nil {
+		return fmt.Errorf("failed to delete user from ourdb: %w", err)
+	}
+
+	return nil
+}
+
+// AccountStore handles the storage and retrieval of Account objects
+type AccountStore struct {
+	Store *DBStore
+}
+
+// NewAccountStore creates a new AccountStore using the provided DBStore
+func NewAccountStore(store *DBStore) *AccountStore {
+	return &AccountStore{
+		Store: store,
+	}
+}
+
+// Save saves an Account to the database
+func (s *AccountStore) Save(account *Account) error {
+	// Validate account data
+	if err := account.Validate(); err != nil {
+		return err
+	}
+
+	// If ID is 0, this is a new account
+	if account.ID == 0 {
+		// Get next ID
+		nextID, err := s.Store.DB.GetNextID()
+		if err != nil {
+			return fmt.Errorf("failed to get next ID: %w", err)
+		}
+		account.ID = nextID
+
+		// Save to ourdb
+		data := account.Serialize()
+		_, err = s.Store.DB.Set(ourdb.OurDBSetArgs{
+			Data: data,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to save account to ourdb: %w", err)
+		}
+
+		// Update user's accounts list
+		userStore := &UserStore{Store: s.Store}
+		user, err := userStore.GetByID(account.UserID)
+		if err != nil {
+			return fmt.Errorf("failed to get user: %w", err)
+		}
+
+		// Add account ID to user's accounts
+		user.Accounts = append(user.Accounts, account.ID)
+
+		// Save updated user
+		err = userStore.Save(user)
+		if err != nil {
+			return fmt.Errorf("failed to update user's accounts: %w", err)
+		}
+	} else {
+		// Update existing account
+		// Get existing account to verify it exists
+		_, err := s.GetByID(account.ID)
+		if err != nil {
+			return fmt.Errorf("failed to get existing account: %w", err)
+		}
+
+		// Save to ourdb
+		data := account.Serialize()
+		id := account.ID
+		_, err = s.Store.DB.Set(ourdb.OurDBSetArgs{
+			ID:   &id,
+			Data: data,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to update account in ourdb: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// GetByID retrieves an Account by ID
+func (s *AccountStore) GetByID(id uint32) (*Account, error) {
+	data, err := s.Store.DB.Get(id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get account from ourdb: %w", err)
+	}
+
+	account, err := DeserializeAccount(data)
+	if err != nil {
+		return nil, fmt.Errorf("failed to deserialize account: %w", err)
+	}
+
+	return account, nil
+}
+
+// GetByUserID retrieves all Accounts for a user
+func (s *AccountStore) GetByUserID(userID uint32) ([]*Account, error) {
+	// Get user to get account IDs
+	userStore := &UserStore{Store: s.Store}
+	user, err := userStore.GetByID(userID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get user: %w", err)
+	}
+
+	// Get accounts
+	accounts := make([]*Account, 0, len(user.Accounts))
+	for _, accountID := range user.Accounts {
+		account, err := s.GetByID(accountID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get account %d: %w", accountID, err)
+		}
+		accounts = append(accounts, account)
+	}
+
+	return accounts, nil
+}
+
+// Delete deletes an Account by ID
+func (s *AccountStore) Delete(id uint32) error {
+	// Get account to get user ID
+	account, err := s.GetByID(id)
+	if err != nil {
+		return fmt.Errorf("failed to get account: %w", err)
+	}
+
+	// Get user to remove account from accounts list
+	userStore := &UserStore{Store: s.Store}
+	user, err := userStore.GetByID(account.UserID)
+	if err != nil {
+		return fmt.Errorf("failed to get user: %w", err)
+	}
+
+	// Remove account ID from user's accounts
+	for i, accountID := range user.Accounts {
+		if accountID == id {
+			user.Accounts = append(user.Accounts[:i], user.Accounts[i+1:]...)
+			break
+		}
+	}
+
+	// Save updated user
+	err = userStore.Save(user)
+	if err != nil {
+		return fmt.Errorf("failed to update user's accounts: %w", err)
+	}
+
+	// Delete from ourdb
+	err = s.Store.DB.Delete(id)
+	if err != nil {
+		return fmt.Errorf("failed to delete account from ourdb: %w", err)
+	}
+
+	return nil
+}
+
+// TransactionStore handles the storage and retrieval of Transaction objects
+type TransactionStore struct {
+	Store *DBStore
+}
+
+// NewTransactionStore creates a new TransactionStore using the provided DBStore
+func NewTransactionStore(store *DBStore) *TransactionStore {
+	return &TransactionStore{
+		Store: store,
+	}
+}
+
+// Save saves a Transaction to the database and updates account balances
+func (s *TransactionStore) Save(transaction *Transaction) error {
+	// Validate transaction data
+	if err := transaction.Validate(); err != nil {
+		return err
+	}
+
+	// Get account store
+	accountStore := &AccountStore{Store: s.Store}
+
+	// Get source account
+	fromAccount, err := accountStore.GetByID(transaction.From)
+	if err != nil {
+		return fmt.Errorf("failed to get source account: %w", err)
+	}
+
+	// Get destination account
+	toAccount, err := accountStore.GetByID(transaction.To)
+	if err != nil {
+		return fmt.Errorf("failed to get destination account: %w", err)
+	}
+
+	// Check if source account has enough balance
+	if fromAccount.Amount < transaction.Amount {
+		return errors.New("insufficient balance in source account")
+	}
+
+	// If TxID is 0, this is a new transaction
+	if transaction.TxID == 0 {
+		// Get next ID
+		nextID, err := s.Store.DB.GetNextID()
+		if err != nil {
+			return fmt.Errorf("failed to get next ID: %w", err)
+		}
+		transaction.TxID = nextID
+
+		// Update account balances
+		fromAccount.Amount -= transaction.Amount
+		toAccount.Amount += transaction.Amount
+
+		// Save updated accounts
+		if err := accountStore.Save(fromAccount); err != nil {
+			return fmt.Errorf("failed to update source account: %w", err)
+		}
+
+		if err := accountStore.Save(toAccount); err != nil {
+			// Rollback source account change if destination update fails
+			fromAccount.Amount += transaction.Amount
+			if rollbackErr := accountStore.Save(fromAccount); rollbackErr != nil {
+				return fmt.Errorf("failed to update destination account and rollback failed: %w, rollback error: %v", err, rollbackErr)
+			}
+			return fmt.Errorf("failed to update destination account: %w", err)
+		}
+
+		// Save transaction to ourdb
+		data := transaction.Serialize()
+		_, err = s.Store.DB.Set(ourdb.OurDBSetArgs{
+			Data: data,
+		})
+		if err != nil {
+			// Rollback account changes if transaction save fails
+			fromAccount.Amount += transaction.Amount
+			toAccount.Amount -= transaction.Amount
+			if rollbackErr1 := accountStore.Save(fromAccount); rollbackErr1 != nil {
+				return fmt.Errorf("failed to save transaction and rollback failed: %w, rollback error 1: %v", err, rollbackErr1)
+			}
+			if rollbackErr2 := accountStore.Save(toAccount); rollbackErr2 != nil {
+				return fmt.Errorf("failed to save transaction and rollback failed: %w, rollback error 2: %v", err, rollbackErr2)
+			}
+			return fmt.Errorf("failed to save transaction to ourdb: %w", err)
+		}
+	} else {
+		// Updating existing transactions is not allowed as it would affect account balances
+		return errors.New("updating existing transactions is not allowed")
+	}
+
+	return nil
+}
+
+// GetByID retrieves a Transaction by ID
+func (s *TransactionStore) GetByID(id uint32) (*Transaction, error) {
+	data, err := s.Store.DB.Get(id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get transaction from ourdb: %w", err)
+	}
+
+	transaction, err := DeserializeTransaction(data)
+	if err != nil {
+		return nil, fmt.Errorf("failed to deserialize transaction: %w", err)
+	}
+
+	return transaction, nil
+}
+
+// GetByAccount retrieves all Transactions for an account (either as source or destination)
+func (s *TransactionStore) GetByAccount(accountID uint32) ([]*Transaction, error) {
+	// This is a simplified implementation that would be inefficient in a real system
+	// In a real system, we would maintain indexes for account transactions
+
+	// Get the next ID to determine the range of transaction IDs
+	nextID, err := s.Store.DB.GetNextID()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get next ID: %w", err)
+	}
+
+	transactions := make([]*Transaction, 0)
+
+	// Iterate through all transactions (this is inefficient but works for a simple implementation)
+	for i := uint32(1); i < nextID; i++ {
+		transaction, err := s.GetByID(i)
+		if err != nil {
+			// Skip if transaction doesn't exist or can't be deserialized
+			continue
+		}
+
+		// Include transaction if it involves the specified account
+		if transaction.From == accountID || transaction.To == accountID {
+			transactions = append(transactions, transaction)
+		}
+	}
+
+	return transactions, nil
+}
+
+// Delete is not implemented for transactions as it would affect account balances
+// In a real system, you might want to implement a "void" or "reverse" transaction instead
+func (s *TransactionStore) Delete(id uint32) error {
+	return errors.New("deleting transactions is not allowed")
+}