moteva/server/internal/modules/auth/memory_store.go

package auth

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

	"img_infinite_canvas/internal/domain/design"
)

type MemoryStore struct {
	mu       sync.RWMutex
	users    map[string]User
	codes    map[string]VerificationCode
	sessions map[string]DeviceSession
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		users:    make(map[string]User),
		codes:    make(map[string]VerificationCode),
		sessions: make(map[string]DeviceSession),
	}
}

func (s *MemoryStore) FindUserByID(ctx context.Context, id string) (User, error) {
	if err := ctx.Err(); err != nil {
		return User{}, err
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	user, ok := s.users[strings.TrimSpace(id)]
	if !ok {
		return User{}, design.ErrNotFound
	}
	return user, nil
}

func (s *MemoryStore) FindUserByPhone(ctx context.Context, countryCode string, phone string) (User, error) {
	return s.find(ctx, func(user User) bool {
		return user.CountryCode == normalizeCountryCode(countryCode) && user.Phone == normalizePhone(phone)
	})
}

func (s *MemoryStore) FindUserByEmail(ctx context.Context, email string) (User, error) {
	email = normalizeEmail(email)
	return s.find(ctx, func(user User) bool {
		return normalizeEmail(user.Email) == email
	})
}

func (s *MemoryStore) FindUserByGoogleSubject(ctx context.Context, subject string) (User, error) {
	subject = strings.TrimSpace(subject)
	return s.find(ctx, func(user User) bool {
		return user.GoogleSubject == subject
	})
}

func (s *MemoryStore) FindUserByWechat(ctx context.Context, openID string, unionID string) (User, error) {
	openID = strings.TrimSpace(openID)
	unionID = strings.TrimSpace(unionID)
	return s.find(ctx, func(user User) bool {
		if unionID != "" && user.WechatUnionID == unionID {
			return true
		}
		return openID != "" && user.WechatOpenID == openID
	})
}

func (s *MemoryStore) UpsertUser(ctx context.Context, user User) (User, error) {
	if err := ctx.Err(); err != nil {
		return User{}, err
	}
	user.ID = normalizeUUID(user.ID)
	user.Email = normalizeEmail(user.Email)
	user.Phone = normalizePhone(user.Phone)
	user.CountryCode = normalizeCountryCode(user.CountryCode)
	user.Name = strings.TrimSpace(user.Name)
	user.AvatarURL = strings.TrimSpace(user.AvatarURL)
	user.Status = activeStatus(user.Status)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.users[user.ID] = user
	return user, nil
}

func (s *MemoryStore) UpdateUserProfile(ctx context.Context, userID string, patch ProfilePatch) (User, error) {
	if err := ctx.Err(); err != nil {
		return User{}, err
	}
	userID = normalizeUUID(userID)
	s.mu.Lock()
	defer s.mu.Unlock()
	user, ok := s.users[userID]
	if !ok {
		return User{}, design.ErrNotFound
	}
	if patch.Name != nil {
		user.Name = strings.TrimSpace(*patch.Name)
	}
	if patch.AvatarURL != nil {
		user.AvatarURL = strings.TrimSpace(*patch.AvatarURL)
	}
	s.users[user.ID] = user
	return user, nil
}

func (s *MemoryStore) CreateDeviceSession(ctx context.Context, session DeviceSession, activeLimit int64) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.sessions[session.ID]; !exists {
		s.sessions[session.ID] = session
	}
	s.enforceDeviceSessionLimitLocked(session.UserID, session.DeviceType, session.ID, activeLimit, session.CreatedAt)
	return nil
}

func (s *MemoryStore) FindDeviceSession(ctx context.Context, id string) (DeviceSession, error) {
	if err := ctx.Err(); err != nil {
		return DeviceSession{}, err
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	session, ok := s.sessions[strings.TrimSpace(id)]
	if !ok {
		return DeviceSession{}, design.ErrNotFound
	}
	return session, nil
}

func (s *MemoryStore) ListDeviceSessions(ctx context.Context, userID string, now time.Time) ([]DeviceSession, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	userID = normalizeUUID(userID)
	s.mu.RLock()
	defer s.mu.RUnlock()
	sessions := make([]DeviceSession, 0)
	for _, session := range s.sessions {
		if session.UserID != userID || session.RevokedAt != nil || !session.ExpiresAt.After(now) {
			continue
		}
		sessions = append(sessions, session)
	}
	return sessions, nil
}

func (s *MemoryStore) TouchDeviceSession(ctx context.Context, session DeviceSession) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	stored, ok := s.sessions[session.ID]
	if !ok || stored.UserID != session.UserID {
		return design.ErrNotFound
	}
	if stored.RevokedAt != nil || !stored.ExpiresAt.After(session.LastSeenAt) {
		return design.ErrNotFound
	}
	if stored.LastSeenAt.After(session.LastSeenAt) {
		return nil
	}
	stored.DeviceType = session.DeviceType
	stored.System = session.System
	stored.Browser = session.Browser
	stored.Client = session.Client
	stored.UserAgent = session.UserAgent
	stored.LastSeenAt = session.LastSeenAt
	s.sessions[session.ID] = stored
	return nil
}

func (s *MemoryStore) EnforceDeviceSessionLimit(
	ctx context.Context,
	userID string,
	deviceType string,
	keepSessionID string,
	activeLimit int64,
	revokedAt time.Time,
) (int64, error) {
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.enforceDeviceSessionLimitLocked(normalizeUUID(userID), deviceType, keepSessionID, activeLimit, revokedAt), nil
}

func (s *MemoryStore) enforceDeviceSessionLimitLocked(
	userID string,
	deviceType string,
	keepSessionID string,
	activeLimit int64,
	revokedAt time.Time,
) int64 {
	if activeLimit < 1 {
		activeLimit = 1
	}
	active := make([]DeviceSession, 0)
	for _, session := range s.sessions {
		if session.UserID != userID || session.DeviceType != deviceType || session.RevokedAt != nil || !session.ExpiresAt.After(revokedAt) {
			continue
		}
		active = append(active, session)
	}
	sort.SliceStable(active, func(i, j int) bool {
		iKeep := active[i].ID == keepSessionID
		jKeep := active[j].ID == keepSessionID
		if iKeep != jKeep {
			return iKeep
		}
		if !active[i].CreatedAt.Equal(active[j].CreatedAt) {
			return active[i].CreatedAt.After(active[j].CreatedAt)
		}
		if !active[i].LastSeenAt.Equal(active[j].LastSeenAt) {
			return active[i].LastSeenAt.After(active[j].LastSeenAt)
		}
		return active[i].ID > active[j].ID
	})

	var removed int64
	for index := int(activeLimit); index < len(active); index++ {
		session := active[index]
		revoked := revokedAt
		session.RevokedAt = &revoked
		s.sessions[session.ID] = session
		removed++
	}
	return removed
}

func (s *MemoryStore) RevokeOtherDeviceSessions(ctx context.Context, userID string, currentSessionID string, revokedAt time.Time) (int64, error) {
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	userID = normalizeUUID(userID)
	s.mu.Lock()
	defer s.mu.Unlock()
	var removed int64
	for id, session := range s.sessions {
		if session.UserID != userID || session.ID == currentSessionID || session.RevokedAt != nil || !session.ExpiresAt.After(revokedAt) {
			continue
		}
		revoked := revokedAt
		session.RevokedAt = &revoked
		s.sessions[id] = session
		removed++
	}
	return removed, nil
}

func (s *MemoryStore) RevokeDeviceSession(ctx context.Context, userID string, sessionID string, revokedAt time.Time) (int64, error) {
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	userID = normalizeUUID(userID)
	s.mu.Lock()
	defer s.mu.Unlock()
	session, ok := s.sessions[sessionID]
	if !ok || session.UserID != userID || session.RevokedAt != nil || !session.ExpiresAt.After(revokedAt) {
		return 0, nil
	}
	revoked := revokedAt
	session.RevokedAt = &revoked
	s.sessions[sessionID] = session
	return 1, nil
}

func (s *MemoryStore) CreateVerificationCode(ctx context.Context, code VerificationCode) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.codes[code.ID] = code
	return nil
}

func (s *MemoryStore) FindActiveVerificationCode(ctx context.Context, region string, channel string, target string, purpose string) (VerificationCode, error) {
	if err := ctx.Err(); err != nil {
		return VerificationCode{}, err
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	var latest VerificationCode
	for _, code := range s.codes {
		if code.Region != region || code.Channel != channel || code.Target != target || code.Purpose != purpose {
			continue
		}
		if code.ConsumedAt != nil {
			continue
		}
		if latest.ID == "" || code.ExpiresAt.After(latest.ExpiresAt) {
			latest = code
		}
	}
	if latest.ID == "" {
		return VerificationCode{}, design.ErrNotFound
	}
	return latest, nil
}

func (s *MemoryStore) MarkVerificationCodeConsumed(ctx context.Context, id string, consumedAt time.Time) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	code, ok := s.codes[id]
	if !ok {
		return design.ErrNotFound
	}
	code.ConsumedAt = &consumedAt
	s.codes[id] = code
	return nil
}

func (s *MemoryStore) IncrementVerificationCodeAttempts(ctx context.Context, id string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	code, ok := s.codes[id]
	if !ok {
		return design.ErrNotFound
	}
	code.Attempts++
	s.codes[id] = code
	return nil
}

func (s *MemoryStore) Close() error {
	return nil
}

func (s *MemoryStore) find(ctx context.Context, match func(User) bool) (User, error) {
	if err := ctx.Err(); err != nil {
		return User{}, err
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, user := range s.users {
		if match(user) {
			return user, nil
		}
	}
	return User{}, design.ErrNotFound
}