package pool

import (
	"net"
	"sync"
	"sync/atomic"
)

var (
	defaultMaxConns = 20
	defaultStrategy = &RoundrobinStrategy{
		lastSelected: initialRoundrobinAtomic(),
	}
)

func initialRoundrobinAtomic() atomic.Int64 {
	a := atomic.Int64{}
	a.Store(-1)
	return a
}

type ConnSelectionStrategy interface {
	Select() int
}

type RoundrobinStrategy struct {
	lastSelected atomic.Int64
}

func (r *RoundrobinStrategy) Select() int {
	return int(r.lastSelected.Add(1))
}

type ConnPool struct {
	m        sync.RWMutex
	strategy ConnSelectionStrategy
	conns    []net.Conn
}

type newConnPoolOpts struct {
	strategy ConnSelectionStrategy
	maxConns int
}

func newNewConnPoolOpts() newConnPoolOpts {
	return newConnPoolOpts{
		strategy: defaultStrategy,
		maxConns: defaultMaxConns,
	}
}

type NewConnPoolOpt func(p *newConnPoolOpts)

func WithStrategy(s ConnSelectionStrategy) NewConnPoolOpt {
	return func(p *newConnPoolOpts) {
		p.strategy = s
	}
}

func WithMaxConns(mc int) NewConnPoolOpt {
	return func(p *newConnPoolOpts) {
		p.maxConns = mc
	}
}

func NewConnPool(opts ...NewConnPoolOpt) *ConnPool {
	o := newNewConnPoolOpts()
	for _, opt := range opts {
		opt(&o)
	}
	return &ConnPool{
		conns:    make([]net.Conn, 0),
		strategy: o.strategy,
	}
}

func (p *ConnPool) SelectConn() net.Conn {
	p.m.RLock()
	defer p.m.RUnlock()
	return p.conns[p.strategy.Select()]
}

func (p *ConnPool) AddConn(conn net.Conn) {
	p.m.Lock()
	defer p.m.Unlock()
	p.conns = append(p.conns, conn)
}