package realtime

import (
	"context"
	"sync"
)

type MemoryBus struct {
	mu          sync.RWMutex
	subscribers map[string]map[*memorySubscription]struct{}
}

type memorySubscription struct {
	bus       *MemoryBus
	projectID string
	ch        chan Event
	once      sync.Once
}

func NewMemoryBus() *MemoryBus {
	return &MemoryBus{subscribers: make(map[string]map[*memorySubscription]struct{})}
}

func (b *MemoryBus) Publish(ctx context.Context, event Event) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	b.mu.RLock()
	subscribers := make([]*memorySubscription, 0, len(b.subscribers[event.ProjectID]))
	for sub := range b.subscribers[event.ProjectID] {
		subscribers = append(subscribers, sub)
	}
	b.mu.RUnlock()
	for _, sub := range subscribers {
		select {
		case sub.ch <- event:
		default:
		}
	}
	return nil
}

func (b *MemoryBus) Subscribe(ctx context.Context, projectID string) (Subscription, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	sub := &memorySubscription{
		bus:       b,
		projectID: projectID,
		ch:        make(chan Event, 32),
	}
	b.mu.Lock()
	if b.subscribers[projectID] == nil {
		b.subscribers[projectID] = make(map[*memorySubscription]struct{})
	}
	b.subscribers[projectID][sub] = struct{}{}
	b.mu.Unlock()
	go func() {
		<-ctx.Done()
		_ = sub.Close()
	}()
	return sub, nil
}

func (b *MemoryBus) Close() error {
	b.mu.Lock()
	defer b.mu.Unlock()
	for _, subscribers := range b.subscribers {
		for sub := range subscribers {
			sub.once.Do(func() {
				sub.closeLocked()
			})
		}
	}
	b.subscribers = make(map[string]map[*memorySubscription]struct{})
	return nil
}

func (s *memorySubscription) C() <-chan Event {
	return s.ch
}

func (s *memorySubscription) Close() error {
	s.once.Do(func() {
		s.bus.mu.Lock()
		defer s.bus.mu.Unlock()
		s.closeLocked()
	})
	return nil
}

func (s *memorySubscription) closeLocked() {
	if subscribers := s.bus.subscribers[s.projectID]; subscribers != nil {
		delete(subscribers, s)
		if len(subscribers) == 0 {
			delete(s.bus.subscribers, s.projectID)
		}
	}
	close(s.ch)
}