package broadcaster

import (
	"encoding/json"
	"sync"
)

type Broadcaster struct {
	chs []chan []byte
	lk  sync.Mutex
}

func NewBroadcaster() *Broadcaster {
	return &Broadcaster{}
}

func (bc *Broadcaster) Register() chan []byte {
	ch := make(chan []byte, 1000)
	bc.lk.Lock()
	bc.chs = append(bc.chs, ch)
	bc.lk.Unlock()
	return ch
}

func (bc *Broadcaster) Unregister(ch chan []byte) {
	bc.lk.Lock()
	defer bc.lk.Unlock()
	for i, c := range bc.chs {
		if c == ch {
			bc.chs = append(bc.chs[:i], bc.chs[i+1:]...)
			close(ch)
			return
		}
	}
}

func (bc *Broadcaster) Broadcast(data map[string]interface{}) {
	b, err := json.Marshal(data)
	if err != nil {
		return
	}
	bc.lk.Lock()
	dead := make([]int, 0)
	for i, ch := range bc.chs {
		select {
		case ch <- b:
		default:
			select {
			case <-ch:
				ch <- b
			default:
				dead = append(dead, i)
			}
		}
	}
	bc.lk.Unlock()
	for j := len(dead) - 1; j >= 0; j-- {
		bc.lk.Lock()
		idx := dead[j]
		if idx < len(bc.chs) {
			close(bc.chs[idx])
			bc.chs = append(bc.chs[:idx], bc.chs[idx+1:]...)
		}
		bc.lk.Unlock()
	}
}