package chain

import "gorm.io/gorm"

type Repository struct {
	db *gorm.DB
}

func NewRepository(db *gorm.DB) *Repository {
	return &Repository{db: db}
}

func (r *Repository) Create(w *UserWallet) error {
	return r.db.Create(w).Error
}

func (r *Repository) FindByUserID(userID uint) (*UserWallet, error) {
	var w UserWallet
	if err := r.db.Where("user_id = ?", userID).First(&w).Error; err != nil {
		return nil, err
	}
	return &w, nil
}

func (r *Repository) FindByPubKeyHex(pubKeyHex string) (*UserWallet, error) {
	var w UserWallet
	if err := r.db.Where("pub_key_hex = ?", pubKeyHex).First(&w).Error; err != nil {
		return nil, err
	}
	return &w, nil
}

// FindAllByKeyVersion returns all wallets with the given key version.
func (r *Repository) FindAllByKeyVersion(version int) ([]UserWallet, error) {
	var wallets []UserWallet
	if err := r.db.Where("key_version = ?", version).Find(&wallets).Error; err != nil {
		return nil, err
	}
	return wallets, nil
}

// UpdateEncryption updates the encryption fields of a wallet.
func (r *Repository) UpdateEncryption(id uint, encPrivKey, encNonce, hkdfSalt string, keyVersion int) error {
	return r.db.Model(&UserWallet{}).Where("id = ?", id).Updates(map[string]any{
		"encrypted_priv_key": encPrivKey,
		"enc_nonce":          encNonce,
		"hkdf_salt":          hkdfSalt,
		"key_version":        keyVersion,
	}).Error
}