package player

import (
	"encoding/base64"
	"math/bits"

	"gopkg.in/yaml.v3"
)

type RoomSet struct {
	bits []byte
}

func (rs *RoomSet) Add(roomID int) {
	if roomID < 0 {
		return
	}
	byteIdx := roomID / 8
	if byteIdx >= len(rs.bits) {
		nb := make([]byte, byteIdx+1)
		copy(nb, rs.bits)
		rs.bits = nb
	}
	rs.bits[byteIdx] |= 1 << (roomID % 8)
}

func (rs *RoomSet) Has(roomID int) bool {
	if roomID < 0 {
		return false
	}
	byteIdx := roomID / 8
	return byteIdx < len(rs.bits) && rs.bits[byteIdx]&(1<<(roomID%8)) != 0
}

func (rs *RoomSet) Len() int {
	n := 0
	for _, b := range rs.bits {
		n += bits.OnesCount8(b)
	}
	return n
}

func (rs RoomSet) IsZero() bool {
	return len(rs.bits) == 0
}

func (rs *RoomSet) Keys() []int {
	var keys []int
	for byteIdx, b := range rs.bits {
		if b == 0 {
			continue
		}
		base := byteIdx * 8
		for bitIdx := 0; bitIdx < 8; bitIdx++ {
			if b&(1<<bitIdx) != 0 {
				keys = append(keys, base+bitIdx)
			}
		}
	}
	return keys
}

func (rs *RoomSet) Rekey(idMap map[int]int) (RoomSet, bool) {
	var next RoomSet
	changed := false
	for _, k := range rs.Keys() {
		newK, ok := idMap[k]
		if !ok {
			newK = k
		} else if newK != k {
			changed = true
		}
		next.Add(newK)
	}
	return next, changed
}

func NewRoomSetFrom(ids ...int) RoomSet {
	var rs RoomSet
	for _, id := range ids {
		rs.Add(id)
	}
	return rs
}

func (rs RoomSet) MarshalYAML() (any, error) {
	return base64.StdEncoding.EncodeToString(rs.bits), nil
}

func (rs *RoomSet) UnmarshalYAML(value *yaml.Node) error {
	if value.Kind == yaml.ScalarNode && (value.Tag == "!!null" || value.Value == "null") {
		rs.bits = nil
		return nil
	}
	var s string
	if err := value.Decode(&s); err != nil {
		return err
	}
	if s == "" {
		rs.bits = nil
		return nil
	}
	b, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return err
	}
	rs.bits = b
	return nil
}