package cache

import (
	"github.com/gorilla/websocket"
)

func NewConnectionCache() *ConnectionCache {
	c := &ConnectionCache{}
	c.state.Store(&connectionMappings{
		id2c: make(map[string]*conn),
	})

	return c
}

func (c *ConnectionCache) Set(id string, connection *websocket.Conn, meta *ConnectionMetaData) {
	c.mu.Lock()
	defer c.mu.Unlock()

	current := c.state.Load()
	next := current.clone()

	next.id2c[id] = &conn{
		connection: connection,
		meta:       meta,
	}

	c.state.Store(next)
}

func (c *ConnectionCache) GetById(id string) (*websocket.Conn, *ConnectionMetaData, bool) {
	m := c.state.Load()
	if conn, ok := m.id2c[id]; ok {
		return conn.connection, conn.meta, true
	}

	return nil, nil, false
}

func (c *ConnectionCache) RemoveById(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	current := c.state.Load()
	if _, ok := current.id2c[id]; !ok {
		return
	}

	next := current.clone()
	delete(next.id2c, id)

	c.state.Store(next)
}

func (c *ConnectionCache) Range(fn func(id string, connection *websocket.Conn, meta *ConnectionMetaData)) {
	m := c.state.Load()
	for id, conn := range m.id2c {
		fn(id, conn.connection, conn.meta)
	}
}

func (c *ConnectionCache) GetAllConnections() []*websocket.Conn {
	m := c.state.Load()
	conns := make([]*websocket.Conn, 0, len(m.id2c))
	for _, conn := range m.id2c {
		conns = append(conns, conn.connection)
	}
	return conns
}

func (c *ConnectionCache) Clear() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.state.Store(&connectionMappings{
		id2c: make(map[string]*conn),
	})
}

func (cm *connectionMappings) clone() *connectionMappings {
	newCM := &connectionMappings{
		id2c: make(map[string]*conn, len(cm.id2c)),
	}
	for k, v := range cm.id2c {
		newCM.id2c[k] = v
	}
	return newCM
}