package main

import (
	"io"
	"io/ioutil"
	"sync"

	"github.com/gorilla/websocket"
)

type Conn struct {
	ws   websocket.Conn
	lock sync.Mutex
}

type Pool struct {
	lock  *sync.RWMutex
	conns *sync.Map //map[string]*Conn
}

func NewPool() *Pool {
	return &Pool{
		conns: &sync.Map{},
		lock:  &sync.RWMutex{},
	}
}

func (p *Pool) Push(id string, conn *websocket.Conn) {
	p.lock.Lock()
	defer p.lock.Unlock()
	p.conns.Store(id, &Conn{ws: *conn})
}

func (p *Pool) Broadcast(mt int, r io.Reader) error {
	p.lock.RLock()
	defer p.lock.RUnlock()
	// io.MultiWriter exists but I like this
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}
	n := 1000000
	cnt := 0
	p.conns.Range(func(k, v interface{}) bool {
		k = k.(string)
		conn := &v.(*Conn).ws
		cnt += 1
		w, err := conn.NextWriter(mt)
		if err != nil {
			p.conns.Delete(k)
			return true
		}
		defer w.Close()
		m, err := w.Write(b)
		if err != nil {
			p.conns.Delete(k)
			return true
		}
		if m < n {
			n = m
		}
		return true
	})
	return nil
}