package node

import (
	"fmt"
	"sync"

	toki "git.toki-labs.com/toki/common-proto-socket/go"

	iop "iop/proto/gen/iop"
)

// NodeEntry represents one connected node.
type NodeEntry struct {
	NodeID   string
	Client   *toki.TcpClient
	Adapters []*iop.AdapterInfo
}

// Registry manages all nodes connected to edge.
type Registry struct {
	mu    sync.RWMutex
	nodes map[string]*NodeEntry
}

func NewRegistry() *Registry {
	return &Registry{nodes: make(map[string]*NodeEntry)}
}

func (r *Registry) Register(entry *NodeEntry) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.nodes[entry.NodeID] = entry
}

func (r *Registry) Unregister(nodeID string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.nodes, nodeID)
}

func (r *Registry) Get(nodeID string) (*NodeEntry, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	e, ok := r.nodes[nodeID]
	return e, ok
}

func (r *Registry) All() []*NodeEntry {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]*NodeEntry, 0, len(r.nodes))
	for _, e := range r.nodes {
		out = append(out, e)
	}
	return out
}

func (r *Registry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.nodes)
}

func (r *Registry) Pick() (*NodeEntry, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, e := range r.nodes {
		return e, nil
	}
	return nil, fmt.Errorf("no nodes connected")
}