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

package model

import (
	"context"
	"sort"
	"sync"
	"time"

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

// MemoryRepo is an in-process store for logic tests.
type MemoryRepo struct {
	mu     sync.Mutex
	wallet map[int64]*entity.Wallet
	signs  map[int64]map[int]*entity.SignRecord
	logs   []*entity.EnergyLog
	nextID int64
}

func NewMemoryRepo() *MemoryRepo {
	return &MemoryRepo{
		wallet: make(map[int64]*entity.Wallet),
		signs:  make(map[int64]map[int]*entity.SignRecord),
		nextID: 1,
	}
}

func (m *MemoryRepo) GetWallet(_ context.Context, userID int64) (*entity.Wallet, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	w, ok := m.wallet[userID]
	if !ok {
		return nil, errs.ErrNotFound
	}
	cp := *w
	return &cp, nil
}

func (m *MemoryRepo) EnsureWallet(_ context.Context, userID int64) (*entity.Wallet, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if w, ok := m.wallet[userID]; ok {
		cp := *w
		return &cp, nil
	}
	w := &entity.Wallet{UserID: userID, Energy: 0, UpdatedMS: time.Now().UnixMilli()}
	m.wallet[userID] = w
	cp := *w
	return &cp, nil
}

func (m *MemoryRepo) AddEnergy(_ context.Context, userID, delta int64, reason string, dayKey int, source string) (*entity.Wallet, *entity.EnergyLog, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	w, ok := m.wallet[userID]
	if !ok {
		w = &entity.Wallet{UserID: userID}
		m.wallet[userID] = w
	}
	w.Energy += delta
	if w.Energy < 0 {
		return nil, nil, errs.ErrInvalidArgument
	}
	w.UpdatedMS = time.Now().UnixMilli()
	id := m.nextID
	m.nextID++
	log := &entity.EnergyLog{
		ID: id, UserID: userID, Delta: delta, BalanceAfter: w.Energy,
		Reason: reason, DayKey: dayKey, Source: source, CreatedMS: w.UpdatedMS,
	}
	m.logs = append(m.logs, log)
	wc, lc := *w, *log
	return &wc, &lc, nil
}

func (m *MemoryRepo) ListSignsInRange(_ context.Context, userID int64, fromDay, toDay int) ([]*entity.SignRecord, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	byDay := m.signs[userID]
	out := make([]*entity.SignRecord, 0)
	for day, rec := range byDay {
		if day >= fromDay && day <= toDay {
			cp := *rec
			out = append(out, &cp)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].DayKey < out[j].DayKey })
	return out, nil
}

func (m *MemoryRepo) GetSign(_ context.Context, userID int64, dayKey int) (*entity.SignRecord, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	byDay := m.signs[userID]
	if byDay == nil {
		return nil, errs.ErrNotFound
	}
	rec, ok := byDay[dayKey]
	if !ok {
		return nil, errs.ErrNotFound
	}
	cp := *rec
	return &cp, nil
}

func (m *MemoryRepo) CreateSign(_ context.Context, rec *entity.SignRecord) error {
	if rec == nil || rec.UserID <= 0 || rec.DayKey <= 0 {
		return errs.ErrInvalidArgument
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	byDay := m.signs[rec.UserID]
	if byDay == nil {
		byDay = make(map[int]*entity.SignRecord)
		m.signs[rec.UserID] = byDay
	}
	if _, ok := byDay[rec.DayKey]; ok {
		return errs.ErrDuplicate
	}
	if rec.ID == 0 {
		rec.ID = m.nextID
		m.nextID++
	}
	if rec.CreatedMS == 0 {
		rec.CreatedMS = time.Now().UnixMilli()
	}
	cp := *rec
	byDay[rec.DayKey] = &cp
	return nil
}

func (m *MemoryRepo) ListEnergyLogs(_ context.Context, userID, cursor int64, limit int) ([]*entity.EnergyLog, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	limit = entity.ClampLimit(limit)
	out := make([]*entity.EnergyLog, 0)
	for i := len(m.logs) - 1; i >= 0; i-- {
		row := m.logs[i]
		if row.UserID != userID {
			continue
		}
		if cursor > 0 && row.ID >= cursor {
			continue
		}
		cp := *row
		out = append(out, &cp)
		if len(out) >= limit {
			break
		}
	}
	return out, nil
}

func (m *MemoryRepo) ListRecentSigns(_ context.Context, userID int64, limit int) ([]*entity.SignRecord, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if limit <= 0 {
		limit = 40
	}
	byDay := m.signs[userID]
	out := make([]*entity.SignRecord, 0, len(byDay))
	for _, rec := range byDay {
		cp := *rec
		out = append(out, &cp)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].DayKey > out[j].DayKey })
	if len(out) > limit {
		out = out[:limit]
	}
	return out, nil
}