package syncz

import (
	"context"
	"sync"
)

type EventCallback[E any] func(ctx context.Context, e E)

type sub[E any] struct {
	done <-chan struct{}
	ch   chan<- E
}

type Subscriptions[E any] struct {
	mu    sync.Mutex
	subs  map[uint64]sub[E]
	index uint64
}

func NewSubscriptions[E any]() *Subscriptions[E] {
	return &Subscriptions[E]{
		subs: make(map[uint64]sub[E]),
	}
}

func (s *Subscriptions[E]) add(sb sub[E]) uint64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	idx := s.index
	s.index++
	s.subs[idx] = sb
	return idx
}

func (s *Subscriptions[E]) remove(idx uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.subs, idx)
}

type Listen[E any] func(cb EventCallback[E])

// Subscribe subscribes to events but, unlike On, does not start listening immediately.
// It returns a function that must be called to listen to events. Make sure the function is always called as it
// performs cleanup once the passed context is done.
// The returned function must only be called once.
func (s *Subscriptions[E]) Subscribe(ctx context.Context) Listen[E] {
	// A small buffer to give consumers a chance to get the event concurrently.
	// We don't want to serialize consumers.
	ch := make(chan E, 1)
	done := ctx.Done()
	sb := sub[E]{
		done: done,
		ch:   ch,
	}
	idx := s.add(sb)
	return func(cb EventCallback[E]) {
		defer s.remove(idx)

		for {
			select {
			case <-done:
				return
			case e := <-ch:
				cb(ctx, e)
			}
		}
	}
}

func (s *Subscriptions[E]) On(ctx context.Context, cb EventCallback[E]) {
	s.Subscribe(ctx)(cb)
}

// Len returns the number of subscriptions.
func (s *Subscriptions[E]) Len() int {
	s.mu.Lock()
	defer s.mu.Unlock()

	return len(s.subs)
}

// Dispatch dispatches the given event to all added subscriptions.
func (s *Subscriptions[E]) Dispatch(ctx context.Context, e E) {
	s.mu.Lock()
	subs := make([]sub[E], 0, len(s.subs))
	for _, sb := range s.subs {
		subs = append(subs, sb)
	}
	s.mu.Unlock()

	done := ctx.Done()
	for _, sb := range subs {
		select {
		case <-done:
			return
		case <-sb.done:
			// It doesn't want the events anymore.
		case sb.ch <- e:
		}
	}
}
