thread-master/apps/backend/internal/module/studio/repository/memory.go

package repository

import (
	"context"
	"sort"
	"sync"

	"apps/backend/internal/module/studio/domain"
)

type MemoryStore struct {
	mu       sync.Mutex
	personas map[string]*domain.Persona
	active   map[int64]string
	plays    map[string]*domain.Play
	outbox   map[string]*domain.OutboxBundle
	ownPosts map[string]*domain.OwnPost
	mentions map[string]*domain.Mention
	syncMeta map[int64]*domain.SyncMeta
}

func NewMemory() *MemoryStore {
	return &MemoryStore{
		personas: map[string]*domain.Persona{},
		active:   map[int64]string{},
		plays:    map[string]*domain.Play{},
		outbox:   map[string]*domain.OutboxBundle{},
		ownPosts: map[string]*domain.OwnPost{},
		mentions: map[string]*domain.Mention{},
		syncMeta: map[int64]*domain.SyncMeta{},
	}
}

func (s *MemoryStore) SavePersona(_ context.Context, p *domain.Persona) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *p
	s.personas[p.ID] = &cp
	return nil
}

func (s *MemoryStore) GetPersona(_ context.Context, id string) (*domain.Persona, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	p, ok := s.personas[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	cp := *p
	return &cp, nil
}

func (s *MemoryStore) ListPersonas(_ context.Context, ownerUID int64) ([]*domain.Persona, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.Persona
	for _, p := range s.personas {
		if p.OwnerUID == ownerUID {
			cp := *p
			out = append(out, &cp)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].UpdatedAt > out[j].UpdatedAt })
	return out, nil
}

func (s *MemoryStore) DeletePersona(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.personas[id]; !ok {
		return domain.ErrNotFound
	}
	delete(s.personas, id)
	return nil
}

func (s *MemoryStore) GetActivePersonaID(_ context.Context, ownerUID int64) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.active[ownerUID], nil
}

func (s *MemoryStore) SetActivePersonaID(_ context.Context, ownerUID int64, personaID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.active[ownerUID] = personaID
	return nil
}

func (s *MemoryStore) SavePlay(_ context.Context, p *domain.Play) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *p
	if p.Steps != nil {
		cp.Steps = append([]domain.PlayStep(nil), p.Steps...)
	}
	if p.CastAccountIDs != nil {
		cp.CastAccountIDs = append([]string(nil), p.CastAccountIDs...)
	}
	if p.TargetExternal != nil {
		te := *p.TargetExternal
		cp.TargetExternal = &te
	}
	s.plays[p.ID] = &cp
	return nil
}

func (s *MemoryStore) GetPlay(_ context.Context, id string) (*domain.Play, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	p, ok := s.plays[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	return copyPlay(p), nil
}

func (s *MemoryStore) ListPlays(_ context.Context, ownerUID int64) ([]*domain.Play, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.Play
	for _, p := range s.plays {
		if p.OwnerUID == ownerUID {
			out = append(out, copyPlay(p))
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].UpdatedAt > out[j].UpdatedAt })
	return out, nil
}

func (s *MemoryStore) DeletePlay(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.plays[id]; !ok {
		return domain.ErrNotFound
	}
	delete(s.plays, id)
	return nil
}

func copyPlay(p *domain.Play) *domain.Play {
	cp := *p
	if p.Steps != nil {
		cp.Steps = append([]domain.PlayStep(nil), p.Steps...)
	}
	if p.CastAccountIDs != nil {
		cp.CastAccountIDs = append([]string(nil), p.CastAccountIDs...)
	}
	if p.TargetExternal != nil {
		te := *p.TargetExternal
		cp.TargetExternal = &te
	}
	return &cp
}

func (s *MemoryStore) SaveOutbox(_ context.Context, b *domain.OutboxBundle) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *b
	if b.Steps != nil {
		cp.Steps = append([]domain.OutboxStep(nil), b.Steps...)
	}
	s.outbox[b.ID] = &cp
	return nil
}

func (s *MemoryStore) GetOutbox(_ context.Context, id string) (*domain.OutboxBundle, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	return copyOutbox(b), nil
}

func (s *MemoryStore) ListOutbox(_ context.Context, ownerUID int64) ([]*domain.OutboxBundle, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.OutboxBundle
	for _, b := range s.outbox {
		if b.OwnerUID == ownerUID {
			out = append(out, copyOutbox(b))
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].UpdatedAt > out[j].UpdatedAt })
	return out, nil
}

func (s *MemoryStore) DeleteOutbox(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[id]
	if !ok {
		return domain.ErrNotFound
	}
	for _, step := range b.Steps {
		if step.Status == domain.StepPublishing {
			return domain.ErrIllegalStatus
		}
	}
	delete(s.outbox, id)
	return nil
}

func (s *MemoryStore) ClaimOutboxStep(_ context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) (*domain.OutboxBundle, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[bundleID]
	if !ok {
		return nil, domain.ErrNotFound
	}
	for i := range b.Steps {
		step := &b.Steps[i]
		if step.ID != stepID {
			continue
		}
		due := step.Status == domain.StepScheduled && step.ScheduledAt <= now
		stale := step.Status == domain.StepPublishing && step.LeaseExpiresAt <= now
		if !due && !stale {
			return nil, domain.ErrIllegalStatus
		}
		step.Status = domain.StepPublishing
		step.Error = ""
		step.LeaseOwner = leaseOwner
		step.LeaseExpiresAt = leaseExpiresAt
		b.Status = domain.OBActive
		b.UpdatedAt = now
		return copyOutbox(b), nil
	}
	return nil, domain.ErrNotFound
}

func (s *MemoryStore) RenewOutboxStepLease(_ context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[bundleID]
	if !ok {
		return domain.ErrNotFound
	}
	for i := range b.Steps {
		step := &b.Steps[i]
		if step.ID == stepID && step.Status == domain.StepPublishing && step.LeaseOwner == leaseOwner && step.LeaseExpiresAt > now {
			step.LeaseExpiresAt = leaseExpiresAt
			b.UpdatedAt = now
			return nil
		}
	}
	return domain.ErrLeaseLost
}

func (s *MemoryStore) FinishOutboxStep(_ context.Context, bundleID, stepID, leaseOwner string, result *domain.OutboxStep, now int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[bundleID]
	if !ok {
		return domain.ErrNotFound
	}
	for i := range b.Steps {
		step := &b.Steps[i]
		if step.ID != stepID || step.Status != domain.StepPublishing || step.LeaseOwner != leaseOwner || step.LeaseExpiresAt <= now {
			continue
		}
		*step = *result
		step.LeaseOwner = ""
		step.LeaseExpiresAt = 0
		if step.Kind == domain.StepRoot {
			for j := range b.Steps {
				reply := &b.Steps[j]
				if reply.Kind != domain.StepReply || reply.Status == domain.StepPublished || reply.Status == domain.StepPublishing {
					continue
				}
				if step.Status == domain.StepFailed {
					reply.Status, reply.Error = domain.StepBlocked, "waiting for root success"
				} else if step.Status == domain.StepPublished && reply.Status == domain.StepBlocked {
					reply.Status, reply.Error = domain.StepScheduled, ""
				}
			}
		}
		b.UpdatedAt = now
		return nil
	}
	return domain.ErrLeaseLost
}

func (s *MemoryStore) RecomputeOutboxStatus(_ context.Context, bundleID string, now int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[bundleID]
	if !ok {
		return domain.ErrNotFound
	}
	b.Status = outboxStatus(b.Steps)
	b.UpdatedAt = now
	return nil
}

func (s *MemoryStore) RetryOutboxStep(_ context.Context, bundleID, stepID string, now int64) (*domain.OutboxBundle, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	b, ok := s.outbox[bundleID]
	if !ok {
		return nil, domain.ErrNotFound
	}
	for i := range b.Steps {
		step := &b.Steps[i]
		if step.ID != stepID {
			continue
		}
		if step.Status != domain.StepFailed {
			return nil, domain.ErrIllegalStatus
		}
		step.Status, step.Error, step.ScheduledAt = domain.StepScheduled, "", now
		step.LeaseOwner, step.LeaseExpiresAt = "", 0
		b.Status, b.UpdatedAt = domain.OBScheduling, now
		return copyOutbox(b), nil
	}
	return nil, domain.ErrNotFound
}

func outboxStatus(steps []domain.OutboxStep) string {
	if len(steps) == 0 {
		return domain.OBCancelled
	}
	allPublished, anyFailed, anyPublishing := true, false, false
	for _, step := range steps {
		allPublished = allPublished && step.Status == domain.StepPublished
		anyFailed = anyFailed || step.Status == domain.StepFailed
		anyPublishing = anyPublishing || step.Status == domain.StepPublishing
	}
	if allPublished {
		return domain.OBCompleted
	}
	if anyFailed {
		return domain.OBPartial
	}
	if anyPublishing {
		return domain.OBActive
	}
	return domain.OBScheduling
}

func (s *MemoryStore) ListAllOutbox(_ context.Context) ([]*domain.OutboxBundle, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.OutboxBundle
	for _, b := range s.outbox {
		out = append(out, copyOutbox(b))
	}
	return out, nil
}

func copyOutbox(b *domain.OutboxBundle) *domain.OutboxBundle {
	cp := *b
	if b.Steps != nil {
		cp.Steps = append([]domain.OutboxStep(nil), b.Steps...)
	}
	return &cp
}

func (s *MemoryStore) SaveOwnPost(_ context.Context, p *domain.OwnPost) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *p
	if p.Replies != nil {
		cp.Replies = append([]domain.OwnPostReply(nil), p.Replies...)
	}
	s.ownPosts[p.ID] = &cp
	return nil
}

func (s *MemoryStore) GetOwnPost(_ context.Context, id string) (*domain.OwnPost, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	p, ok := s.ownPosts[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	cp := *p
	if p.Replies != nil {
		cp.Replies = append([]domain.OwnPostReply(nil), p.Replies...)
	}
	return &cp, nil
}

func (s *MemoryStore) DeleteOwnPost(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.ownPosts, id)
	return nil
}

func (s *MemoryStore) ListOwnPosts(_ context.Context, ownerUID int64, accountID string) ([]*domain.OwnPost, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.OwnPost
	for _, p := range s.ownPosts {
		if p.OwnerUID != ownerUID {
			continue
		}
		if accountID != "" && p.AccountID != accountID {
			continue
		}
		cp := *p
		if p.Replies != nil {
			cp.Replies = append([]domain.OwnPostReply(nil), p.Replies...)
		}
		out = append(out, &cp)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].PublishedAt > out[j].PublishedAt })
	return out, nil
}

func (s *MemoryStore) GetSyncMeta(_ context.Context, ownerUID int64) (*domain.SyncMeta, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	m, ok := s.syncMeta[ownerUID]
	if !ok {
		return &domain.SyncMeta{OwnerUID: ownerUID}, nil
	}
	cp := *m
	return &cp, nil
}

func (s *MemoryStore) SetSyncMeta(_ context.Context, m *domain.SyncMeta) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *m
	s.syncMeta[m.OwnerUID] = &cp
	return nil
}

func (s *MemoryStore) SaveMention(_ context.Context, m *domain.Mention) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *m
	s.mentions[m.ID] = &cp
	return nil
}

func (s *MemoryStore) GetMention(_ context.Context, id string) (*domain.Mention, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	m, ok := s.mentions[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	cp := *m
	return &cp, nil
}

func (s *MemoryStore) ListMentions(_ context.Context, ownerUID int64, accountID string) ([]*domain.Mention, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.Mention
	for _, m := range s.mentions {
		if m.OwnerUID != ownerUID {
			continue
		}
		if accountID != "" && m.AccountID != accountID {
			continue
		}
		cp := *m
		out = append(out, &cp)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
	return out, nil
}

var _ domain.Repository = (*MemoryStore)(nil)