package nostr

import (
	"context"
	"sync"
	"time"

	gn "github.com/nbd-wtf/go-nostr"
)

// Pool is a small relay-connection pool.
type Pool struct {
	mu      sync.Mutex
	relays  []string
	active  map[string]*gn.Relay
}

func NewPool(urls []string) *Pool {
	return &Pool{relays: urls, active: map[string]*gn.Relay{}}
}

func (p *Pool) URLs() []string { return p.relays }

func (p *Pool) Connect(ctx context.Context, url string) (*gn.Relay, error) {
	p.mu.Lock()
	if r, ok := p.active[url]; ok && r.IsConnected() {
		p.mu.Unlock()
		return r, nil
	}
	p.mu.Unlock()

	dialCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	r, err := gn.RelayConnect(dialCtx, url)
	if err != nil {
		return nil, err
	}
	p.mu.Lock()
	p.active[url] = r
	p.mu.Unlock()
	return r, nil
}

func (p *Pool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, r := range p.active {
		_ = r.Close()
	}
	p.active = map[string]*gn.Relay{}
}