backend/internal/store/memory.go

package store

import (
	"errors"
	"sync"
	"time"
)

var (
	ErrNotFound      = errors.New("not found")
	ErrAlreadyExists = errors.New("already exists")
)

type MemoryStore struct {
	mu sync.RWMutex

	users       map[string]User
	challenges  map[string]Challenge
	sessions    map[string]Session
	invitations map[string]Invitation
	collections map[string]Collection
	features    map[string]Feature
	assets      map[string]Asset
	featureRefs map[string]map[string]FeatureAsset
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		users:       make(map[string]User),
		challenges:  make(map[string]Challenge),
		sessions:    make(map[string]Session),
		invitations: make(map[string]Invitation),
		collections: make(map[string]Collection),
		features:    make(map[string]Feature),
		assets:      make(map[string]Asset),
		featureRefs: make(map[string]map[string]FeatureAsset),
	}
}

func (s *MemoryStore) UpsertUser(user User) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.users[user.PublicKey] = user
}

func (s *MemoryStore) GetUser(publicKey string) (User, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	user, ok := s.users[publicKey]
	if !ok {
		return User{}, ErrNotFound
	}
	return user, nil
}

func (s *MemoryStore) CreateChallenge(ch Challenge) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.challenges[ch.Nonce]; exists {
		return ErrAlreadyExists
	}
	s.challenges[ch.Nonce] = ch
	return nil
}

func (s *MemoryStore) GetChallenge(nonce string) (Challenge, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	ch, ok := s.challenges[nonce]
	if !ok {
		return Challenge{}, ErrNotFound
	}
	return ch, nil
}

func (s *MemoryStore) MarkChallengeUsed(nonce string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	ch, ok := s.challenges[nonce]
	if !ok {
		return ErrNotFound
	}
	ch.Used = true
	s.challenges[nonce] = ch
	return nil
}

func (s *MemoryStore) SaveSession(session Session) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[session.Token] = session
}

func (s *MemoryStore) GetSession(token string) (Session, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	session, ok := s.sessions[token]
	if !ok {
		return Session{}, ErrNotFound
	}
	return session, nil
}

func (s *MemoryStore) SaveInvitation(inv Invitation) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.invitations[inv.JTI]; exists {
		return ErrAlreadyExists
	}
	s.invitations[inv.JTI] = inv
	return nil
}

func (s *MemoryStore) GetInvitation(jti string) (Invitation, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	inv, ok := s.invitations[jti]
	if !ok {
		return Invitation{}, ErrNotFound
	}
	return inv, nil
}

func (s *MemoryStore) IncrementInvitationUsage(jti string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	inv, ok := s.invitations[jti]
	if !ok {
		return ErrNotFound
	}
	inv.UsedCount++
	s.invitations[jti] = inv
	return nil
}

func (s *MemoryStore) SaveCollection(c Collection) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.collections[c.ID] = c
}

func (s *MemoryStore) ListCollectionsByOwner(owner string) []Collection {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]Collection, 0)
	for _, c := range s.collections {
		if c.OwnerKey == owner {
			result = append(result, c)
		}
	}
	return result
}

func (s *MemoryStore) GetCollection(id string) (Collection, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	c, ok := s.collections[id]
	if !ok {
		return Collection{}, ErrNotFound
	}
	return c, nil
}

func (s *MemoryStore) DeleteCollection(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.collections[id]; !ok {
		return ErrNotFound
	}
	for fid, f := range s.features {
		if f.CollectionID == id {
			delete(s.features, fid)
			delete(s.featureRefs, fid)
		}
	}
	delete(s.collections, id)
	return nil
}

func (s *MemoryStore) SaveFeature(f Feature) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.features[f.ID] = f
}

func (s *MemoryStore) ListFeaturesByCollection(collectionID string) []Feature {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]Feature, 0)
	for _, f := range s.features {
		if f.CollectionID == collectionID {
			result = append(result, f)
		}
	}
	return result
}

func (s *MemoryStore) GetFeature(featureID string) (Feature, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	f, ok := s.features[featureID]
	if !ok {
		return Feature{}, ErrNotFound
	}
	return f, nil
}

func (s *MemoryStore) DeleteFeature(featureID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.features[featureID]; !ok {
		return ErrNotFound
	}
	delete(s.features, featureID)
	delete(s.featureRefs, featureID)
	return nil
}

func (s *MemoryStore) SaveAsset(a Asset) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.assets[a.ID] = a
}

func (s *MemoryStore) GetAsset(assetID string) (Asset, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	a, ok := s.assets[assetID]
	if !ok {
		return Asset{}, ErrNotFound
	}
	return a, nil
}

func (s *MemoryStore) GetAssetByOwnerChecksumExt(ownerKey, checksum, ext string) (Asset, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, a := range s.assets {
		if a.OwnerKey == ownerKey && a.Checksum == checksum && a.Ext == ext {
			return a, nil
		}
	}
	return Asset{}, ErrNotFound
}

func (s *MemoryStore) SetAssetPublic(assetID string, isPublic bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	a, ok := s.assets[assetID]
	if !ok {
		return ErrNotFound
	}
	a.IsPublic = isPublic
	a.UpdatedAt = time.Now().UTC()
	s.assets[assetID] = a
	return nil
}

func (s *MemoryStore) LinkAssetToFeature(featureID, assetID, name, description string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.features[featureID]; !ok {
		return ErrNotFound
	}
	a, ok := s.assets[assetID]
	if !ok {
		return ErrNotFound
	}
	if _, ok := s.featureRefs[featureID]; !ok {
		s.featureRefs[featureID] = make(map[string]FeatureAsset)
	}
	if existing, exists := s.featureRefs[featureID][assetID]; exists {
		existing.Name = name
		existing.Description = description
		s.featureRefs[featureID][assetID] = existing
		return nil
	}
	s.featureRefs[featureID][assetID] = FeatureAsset{
		Asset:       a,
		FeatureID:   featureID,
		Name:        name,
		Description: description,
		LinkedAt:    time.Now().UTC(),
	}
	return nil
}

func (s *MemoryStore) UnlinkAssetFromFeature(featureID, assetID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	links, ok := s.featureRefs[featureID]
	if !ok {
		return ErrNotFound
	}
	if _, exists := links[assetID]; !exists {
		return ErrNotFound
	}
	delete(links, assetID)
	return nil
}

func (s *MemoryStore) ListAssetsByFeature(featureID string) []FeatureAsset {
	s.mu.RLock()
	defer s.mu.RUnlock()
	links, ok := s.featureRefs[featureID]
	if !ok {
		return []FeatureAsset{}
	}
	result := make([]FeatureAsset, 0, len(links))
	for assetID, fa := range links {
		if updated, exists := s.assets[assetID]; exists {
			fa.Asset = updated
		}
		result = append(result, fa)
	}
	return result
}

func (s *MemoryStore) PruneExpired(now time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for nonce, ch := range s.challenges {
		if ch.ExpiresAt.Before(now) {
			delete(s.challenges, nonce)
		}
	}
	for token, sess := range s.sessions {
		if sess.ExpiresAt.Before(now) {
			delete(s.sessions, token)
		}
	}
	for jti, inv := range s.invitations {
		if inv.ExpiresAt.Before(now) {
			delete(s.invitations, jti)
		}
	}
}