spark-server/internal/module/energy/model/repo.go

package model

import (
	"context"
	"errors"
	"time"

	mysqldriver "github.com/go-sql-driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"spark-server/internal/module/energy/entity"
	"spark-server/internal/pkg/errs"
)

// Repo is MySQL persistence for wallets, sign records, and energy logs.
type Repo struct {
	gdb *gorm.DB
}

func NewRepo(gdb *gorm.DB) *Repo { return &Repo{gdb: gdb} }

func (m *Repo) session(ctx context.Context) (*gorm.DB, error) {
	if m == nil || m.gdb == nil {
		return nil, errs.ErrAvailable
	}
	return m.gdb.WithContext(ctx), nil
}

func mapDBErr(err error) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return errs.ErrNotFound
	}
	var me *mysqldriver.MySQLError
	if errors.As(err, &me) && me.Number == 1062 {
		return errs.ErrDuplicate
	}
	return errs.ErrInner.WithError("energy-mysql", err)
}

func (m *Repo) GetWallet(ctx context.Context, userID int64) (*entity.Wallet, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	var row WalletRow
	if err := db.Where("user_id = ?", userID).Take(&row).Error; err != nil {
		return nil, mapDBErr(err)
	}
	return row.Entity(), nil
}

func (m *Repo) EnsureWallet(ctx context.Context, userID int64) (*entity.Wallet, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	now := time.Now()
	row := WalletRow{UserID: userID, Energy: 0, UpdatedAt: now}
	err = db.Clauses(clause.OnConflict{DoNothing: true}).Create(&row).Error
	if err != nil {
		return nil, mapDBErr(err)
	}
	return m.GetWallet(ctx, userID)
}

func (m *Repo) AddEnergy(ctx context.Context, userID, delta int64, reason string, dayKey int, source string) (*entity.Wallet, *entity.EnergyLog, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, nil, err
	}
	var wallet *entity.Wallet
	var logRow *entity.EnergyLog
	err = db.Transaction(func(tx *gorm.DB) error {
		now := time.Now()
		var w WalletRow
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("user_id = ?", userID).Take(&w).Error; err != nil {
			if !errors.Is(err, gorm.ErrRecordNotFound) {
				return err
			}
			w = WalletRow{UserID: userID, Energy: 0, UpdatedAt: now}
			if err := tx.Create(&w).Error; err != nil {
				return err
			}
		}
		w.Energy += delta
		if w.Energy < 0 {
			return errs.ErrInvalidArgument
		}
		w.UpdatedAt = now
		if err := tx.Save(&w).Error; err != nil {
			return err
		}
		lr := EnergyLogRow{
			UserID: userID, Delta: delta, BalanceAfter: w.Energy, Reason: reason,
			DayKey: dayKey, Source: source, CreatedAt: now,
		}
		if err := tx.Create(&lr).Error; err != nil {
			return err
		}
		wallet = w.Entity()
		logRow = lr.Entity()
		return nil
	})
	if err != nil {
		return nil, nil, mapDBErr(err)
	}
	return wallet, logRow, nil
}

func (m *Repo) ListSignsInRange(ctx context.Context, userID int64, fromDay, toDay int) ([]*entity.SignRecord, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	var rows []SignRow
	if err := db.Where("user_id = ? AND day_key >= ? AND day_key <= ?", userID, fromDay, toDay).
		Order("day_key ASC").Find(&rows).Error; err != nil {
		return nil, mapDBErr(err)
	}
	out := make([]*entity.SignRecord, 0, len(rows))
	for i := range rows {
		out = append(out, rows[i].Entity())
	}
	return out, nil
}

func (m *Repo) GetSign(ctx context.Context, userID int64, dayKey int) (*entity.SignRecord, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	var row SignRow
	if err := db.Where("user_id = ? AND day_key = ?", userID, dayKey).Take(&row).Error; err != nil {
		return nil, mapDBErr(err)
	}
	return row.Entity(), nil
}

func (m *Repo) CreateSign(ctx context.Context, rec *entity.SignRecord) error {
	if rec == nil || rec.UserID <= 0 || rec.DayKey <= 0 {
		return errs.ErrInvalidArgument
	}
	db, err := m.session(ctx)
	if err != nil {
		return err
	}
	row := signFromEntity(rec)
	if row.CreatedAt.IsZero() {
		row.CreatedAt = time.Now()
	}
	return mapDBErr(db.Create(row).Error)
}

func (m *Repo) ListEnergyLogs(ctx context.Context, userID, cursor int64, limit int) ([]*entity.EnergyLog, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	limit = entity.ClampLimit(limit)
	q := db.Where("user_id = ?", userID)
	if cursor > 0 {
		q = q.Where("id < ?", cursor)
	}
	var rows []EnergyLogRow
	if err := q.Order("id DESC").Limit(limit).Find(&rows).Error; err != nil {
		return nil, mapDBErr(err)
	}
	out := make([]*entity.EnergyLog, 0, len(rows))
	for i := range rows {
		out = append(out, rows[i].Entity())
	}
	return out, nil
}

func (m *Repo) ListRecentSigns(ctx context.Context, userID int64, limit int) ([]*entity.SignRecord, error) {
	db, err := m.session(ctx)
	if err != nil {
		return nil, err
	}
	if limit <= 0 {
		limit = 40
	}
	var rows []SignRow
	if err := db.Where("user_id = ?", userID).Order("day_key DESC").Limit(limit).Find(&rows).Error; err != nil {
		return nil, mapDBErr(err)
	}
	out := make([]*entity.SignRecord, 0, len(rows))
	for i := range rows {
		out = append(out, rows[i].Entity())
	}
	return out, nil
}