wallet/pkg/store/account.go

package store

import (
	"fmt"

	"gorm.io/gorm"
)

type NetworkType string

const (
	Tron     NetworkType = "tron"
	Ethereum NetworkType = "ethereum"
)

type Account interface {
	GetAccountID() string
	GetHDIndex() int
	GetAddress() string
	GetSpender() string
	GetLabel() string
}

func CreateAccount(a Account) error {
	db, err := Init()
	if err != nil {
		return err
	}
	if a.GetAccountID() == "" {
		return fmt.Errorf("account_id is required")
	}
	return db.Create(a).Error
}

func UpdateAccount(a Account) error {
	db, err := Init()
	if err != nil {
		return err
	}
	return db.Save(a).Error
}

func GetAccount(accountId string, network NetworkType) (Account, error) {
	db, err := Init()
	if err != nil {
		return nil, err
	}
	var tx *gorm.DB = db.Where("account_id = ?", accountId)
	var acc Account
	switch network {
	case Tron:
		var account TronAccount
		acc = &account
		err = tx.First(&account).Error
	case Ethereum:
		var account EthereumAccount
		acc = &account
		err = tx.First(&account).Error
	default:
		return nil, fmt.Errorf("unknown network type: %s", network)
	}
	if err != nil {
		return nil, err
	}
	return acc, nil
}

func GetAccountByAddress(address string, network NetworkType) (Account, error) {
	db, err := Init()
	if err != nil {
		return nil, err
	}
	var tx *gorm.DB = db.Where("address = ?", address)
	var acc Account
	switch network {
	case Tron:
		var account TronAccount
		acc = &account
		err = tx.First(&account).Error
	case Ethereum:
		var account EthereumAccount
		acc = &account
		err = tx.First(&account).Error
	default:
		return nil, fmt.Errorf("unknown network type: %s", network)
	}
	if err != nil {
		return nil, err
	}
	return acc, nil
}

func CountAccounts(network NetworkType) (int64, error) {
	db, err := Init()
	if err != nil {
		return 0, err
	}
	var count int64
	switch network {
	case Tron:
		var account TronAccount
		err = db.Model(&account).Count(&count).Error
	case Ethereum:
		var account EthereumAccount
		err = db.Model(&account).Count(&count).Error
	default:
		return 0, fmt.Errorf("unknown network type: %s", network)
	}
	if err != nil {
		return 0, err
	}
	return count, nil
}

func LastAccount(network NetworkType) (Account, error) {
	db, err := Init()
	if err != nil {
		return nil, err
	}
	var tx *gorm.DB
	var acc Account

	switch network {
	case Tron:
		var account TronAccount
		acc = &account
		tx = db.Last(&account)
	case Ethereum:
		var account EthereumAccount
		acc = &account
		tx = db.Last(&account)
	default:
		return nil, fmt.Errorf("unknown network type: %s", network)
	}
	if err := tx.Error; err != nil {
		return nil, err
	}
	return acc, nil
}

func ListAccounts(network NetworkType, offset int, limit int) ([]Account, error) {
	db, err := Init()
	if err != nil {
		return nil, err
	}
	var tx *gorm.DB
	switch network {
	case Tron:
		var accounts []TronAccount = []TronAccount{}
		tx = db.Offset(offset).Limit(limit).Find(&accounts)
		if err := tx.Error; err != nil {
			return nil, err
		}
		var convertedAccounts []Account = make([]Account, len(accounts))
		for i := range accounts {
			convertedAccounts[i] = Account(&accounts[i])
		}
		return convertedAccounts, nil
	case Ethereum:
		var accounts []EthereumAccount = []EthereumAccount{}
		tx = db.Offset(offset).Limit(limit).Find(&accounts)
		if err := tx.Error; err != nil {
			return nil, err
		}
		var convertedAccounts []Account = make([]Account, len(accounts))
		for i := range accounts {
			convertedAccounts[i] = Account(&accounts[i])
		}
		return convertedAccounts, nil
	default:
		return nil, fmt.Errorf("unknown network type: %s", network)
	}
}