package sse

import (
	"sync"
)

type Event struct {
	Type string
	Data any
}

type subscriber struct {
	ch     chan Event
	closed bool
}

type Subscription struct {
	hub      *Hub
	key      string
	personal bool
	sub      *subscriber
	once     sync.Once
}

func (s *Subscription) Events() <-chan Event { return s.sub.ch }

func (s *Subscription) Close() {
	s.once.Do(func() {
		if s.personal {
			s.hub.unsubscribeUser(s.key, s.sub)
		} else {
			s.hub.unsubscribeAlbum(s.key, s.sub)
		}
	})
}

type Hub struct {
	mu     sync.RWMutex
	albums map[string]map[*subscriber]struct{}
	users  map[string]map[*subscriber]struct{}
	buffer int
}

func NewHub(buffer int) *Hub {
	if buffer <= 0 {
		buffer = 64
	}
	return &Hub{
		albums: make(map[string]map[*subscriber]struct{}),
		users:  make(map[string]map[*subscriber]struct{}),
		buffer: buffer,
	}
}

func (h *Hub) SubscribeAlbum(albumID string) *Subscription {
	sub := &subscriber{ch: make(chan Event, h.buffer)}
	h.mu.Lock()
	if h.albums[albumID] == nil {
		h.albums[albumID] = make(map[*subscriber]struct{})
	}
	h.albums[albumID][sub] = struct{}{}
	h.mu.Unlock()
	return &Subscription{hub: h, key: albumID, sub: sub}
}

func (h *Hub) SubscribeUser(userID string) *Subscription {
	sub := &subscriber{ch: make(chan Event, h.buffer)}
	h.mu.Lock()
	if h.users[userID] == nil {
		h.users[userID] = make(map[*subscriber]struct{})
	}
	h.users[userID][sub] = struct{}{}
	h.mu.Unlock()
	return &Subscription{hub: h, key: userID, personal: true, sub: sub}
}

func (h *Hub) Publish(albumID, eventType string, data any) {
	if albumID == "" {
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	h.broadcast(h.albums[albumID], Event{Type: eventType, Data: data})
}

func (h *Hub) PublishUser(userID, eventType string, data any) {
	if userID == "" {
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	h.broadcast(h.users[userID], Event{Type: eventType, Data: data})
}

func (h *Hub) broadcast(subs map[*subscriber]struct{}, ev Event) {
	for sub := range subs {
		if sub.closed {
			continue
		}
		select {
		case sub.ch <- ev:
		default:
			// slow client: drop it (backpressure)
			sub.closed = true
			close(sub.ch)
		}
	}
}

func (h *Hub) unsubscribeAlbum(albumID string, sub *subscriber) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if subs, ok := h.albums[albumID]; ok {
		delete(subs, sub)
		if len(subs) == 0 {
			delete(h.albums, albumID)
		}
	}
	closeIfOpen(sub)
}

func (h *Hub) unsubscribeUser(userID string, sub *subscriber) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if subs, ok := h.users[userID]; ok {
		delete(subs, sub)
		if len(subs) == 0 {
			delete(h.users, userID)
		}
	}
	closeIfOpen(sub)
}

func closeIfOpen(sub *subscriber) {
	if !sub.closed {
		sub.closed = true
		close(sub.ch)
	}
}