package node

import (
	"fmt"
	"sync"

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

// NodeEntry represents one connected node.
type NodeEntry struct {
	NodeID string
	Alias  string
	Client *toki.TcpClient
}

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

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

func (r *Registry) Register(entry *NodeEntry) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.byID[entry.NodeID] = entry
	if entry.Alias != "" {
		r.byAlias[entry.Alias] = entry
	}
}

func (r *Registry) Unregister(nodeID string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if entry, ok := r.byID[nodeID]; ok {
		if entry.Alias != "" {
			delete(r.byAlias, entry.Alias)
		}
		delete(r.byID, nodeID)
	}
}

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

func (r *Registry) Resolve(ref string) (*NodeEntry, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if ref != "" {
		if entry, ok := r.byID[ref]; ok {
			return entry, nil
		}
		if entry, ok := r.byAlias[ref]; ok {
			return entry, nil
		}
		return nil, fmt.Errorf("node %q not found", ref)
	}

	if len(r.byID) == 1 {
		for _, entry := range r.byID {
			return entry, nil
		}
	}
	if len(r.byID) == 0 {
		return nil, fmt.Errorf("no nodes connected")
	}
	return nil, fmt.Errorf("multiple nodes connected; select one with /node <id|alias>")
}

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

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

// Pick is deprecated: use Resolve("") instead for single-node fallback, or Resolve(ref) for explicit selection.
func (r *Registry) Pick() (*NodeEntry, error) {
	return r.Resolve("")
}