iop/apps/edge/internal/node/registry.go

package node

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"sync"

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

	"iop/packages/go/config"
)

const (
	LifecycleAccepted  = "accepted"
	LifecycleOnline    = "online"
	LifecycleFailed    = "failed"
	LifecycleConnected = "connected"
)

// NodeEntry represents one connected node.
type NodeEntry struct {
	NodeID         string
	Alias          string
	AgentKind      string
	LifecycleState string
	Client         *toki.TcpClient
	Index          int
	HasIndex       bool
}

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

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

func (r *Registry) Register(entry *NodeEntry) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.registerLocked(entry)
}

// RegisterIfAbsent registers entry only when the node id is not already
// connected. The check and insert happen under one lock so concurrent duplicate
// registration attempts cannot both be accepted by the transport server.
func (r *Registry) RegisterIfAbsent(entry *NodeEntry) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, exists := r.byID[entry.NodeID]; exists {
		return false
	}
	r.registerLocked(entry)
	return true
}

func (r *Registry) registerLocked(entry *NodeEntry) {
	if entry.AgentKind == "" {
		entry.AgentKind = config.AgentKindGenericNode
	}
	if entry.LifecycleState == "" {
		entry.LifecycleState = LifecycleConnected
	}
	if idx, ok := r.indexByID[entry.NodeID]; ok {
		entry.Index = idx
		entry.HasIndex = true
	} else {
		if !entry.HasIndex {
			entry.Index = r.nextIndex
		}
		entry.HasIndex = true
		r.indexByID[entry.NodeID] = entry.Index
		if entry.Index >= r.nextIndex {
			r.nextIndex = entry.Index + 1
		}
	}
	r.byID[entry.NodeID] = entry
	r.byIndex[entry.Index] = entry
	if entry.Alias != "" {
		r.byAlias[entry.Alias] = entry
	}
}

func (r *Registry) UpdateLifecycle(nodeID string, state string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	entry, ok := r.byID[nodeID]
	if !ok {
		return false
	}
	entry.LifecycleState = state
	return true
}

func (e *NodeEntry) DisplayLabel() string {
	if e == nil {
		return "unknown"
	}
	if e.HasIndex {
		return fmt.Sprintf("node%d", e.Index)
	}
	if e.Alias != "" {
		return e.Alias
	}
	if e.NodeID != "" {
		return e.NodeID
	}
	return "unknown"
}

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

// UnregisterIfClient removes nodeID only when the currently registered entry
// belongs to client. Late disconnect callbacks from rejected or superseded
// connections must not clear the live registry entry.
func (r *Registry) UnregisterIfClient(nodeID string, client *toki.TcpClient) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	entry, ok := r.byID[nodeID]
	if !ok || entry.Client != client {
		return false
	}
	r.unregisterLocked(nodeID, entry)
	return true
}

func (r *Registry) unregisterLocked(nodeID string, entry *NodeEntry) {
	if entry.Alias != "" {
		delete(r.byAlias, entry.Alias)
	}
	delete(r.byIndex, entry.Index)
	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
		}
		if entry, ok := r.resolveDisplayLabelLocked(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) resolveDisplayLabelLocked(ref string) (*NodeEntry, bool) {
	idx, ok := parseDisplayNodeIndex(ref)
	if !ok {
		return nil, false
	}
	entry, ok := r.byIndex[idx]
	return entry, ok
}

func parseDisplayNodeIndex(ref string) (int, bool) {
	if !strings.HasPrefix(ref, "node") {
		return 0, false
	}
	raw := strings.TrimPrefix(ref, "node")
	if raw == "" {
		return 0, false
	}
	idx, err := strconv.Atoi(raw)
	if err != nil || idx < 0 {
		return 0, false
	}
	return idx, true
}

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)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].HasIndex != out[j].HasIndex {
			return out[i].HasIndex
		}
		if out[i].HasIndex && out[i].Index != out[j].Index {
			return out[i].Index < out[j].Index
		}
		return out[i].NodeID < out[j].NodeID
	})
	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("")
}