bzzz/pkg/protocol/resolver.go

package protocol

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/peerstore"
)

// PeerCapability represents the capabilities of a peer
type PeerCapability struct {
	PeerID         peer.ID           `json:"peer_id"`
	Agent          string            `json:"agent"`
	Role           string            `json:"role"`
	Capabilities   []string          `json:"capabilities"`
	Models         []string          `json:"models"`
	Specialization string            `json:"specialization"`
	Project        string            `json:"project"`
	LastSeen       time.Time         `json:"last_seen"`
	Status         string            `json:"status"` // "online", "busy", "offline"
	Metadata       map[string]string `json:"metadata"`
}

// PeerAddress represents a resolved peer address
type PeerAddress struct {
	PeerID    peer.ID                `json:"peer_id"`
	Addresses []string               `json:"addresses"`
	Priority  int                    `json:"priority"`
	Metadata  map[string]interface{} `json:"metadata"`
}

// ResolutionResult represents the result of address resolution
type ResolutionResult struct {
	URI           *BzzzURI       `json:"uri"`
	Peers         []*PeerAddress `json:"peers"`
	ResolvedAt    time.Time      `json:"resolved_at"`
	ResolutionTTL time.Duration  `json:"ttl"`
	Strategy      string         `json:"strategy"`
}

// ResolutionStrategy defines how to resolve addresses
type ResolutionStrategy string

const (
	StrategyExact     ResolutionStrategy = "exact"     // Exact match only
	StrategyBestMatch ResolutionStrategy = "best_match" // Best available match
	StrategyLoadBalance ResolutionStrategy = "load_balance" // Load balance among matches
	StrategyPriority  ResolutionStrategy = "priority"   // Highest priority first
)

// Resolver handles semantic address resolution
type Resolver struct {
	// Peer capability registry
	capabilities map[peer.ID]*PeerCapability
	capMutex     sync.RWMutex
	
	// Address resolution cache
	cache    map[string]*ResolutionResult
	cacheMutex sync.RWMutex
	cacheTTL   time.Duration
	
	// Configuration
	defaultStrategy ResolutionStrategy
	maxPeersPerResult int
	
	// Peerstore for address information
	peerstore peerstore.Peerstore
}

// NewResolver creates a new semantic address resolver
func NewResolver(peerstore peerstore.Peerstore, opts ...ResolverOption) *Resolver {
	r := &Resolver{
		capabilities:      make(map[peer.ID]*PeerCapability),
		cache:            make(map[string]*ResolutionResult),
		cacheTTL:         5 * time.Minute,
		defaultStrategy:  StrategyBestMatch,
		maxPeersPerResult: 5,
		peerstore:        peerstore,
	}
	
	for _, opt := range opts {
		opt(r)
	}
	
	// Start background cleanup
	go r.startCleanup()
	
	return r
}

// ResolverOption configures the resolver
type ResolverOption func(*Resolver)

// WithCacheTTL sets the cache TTL
func WithCacheTTL(ttl time.Duration) ResolverOption {
	return func(r *Resolver) {
		r.cacheTTL = ttl
	}
}

// WithDefaultStrategy sets the default resolution strategy
func WithDefaultStrategy(strategy ResolutionStrategy) ResolverOption {
	return func(r *Resolver) {
		r.defaultStrategy = strategy
	}
}

// WithMaxPeersPerResult sets the maximum peers per result
func WithMaxPeersPerResult(max int) ResolverOption {
	return func(r *Resolver) {
		r.maxPeersPerResult = max
	}
}

// RegisterPeer registers a peer's capabilities
func (r *Resolver) RegisterPeer(peerID peer.ID, capability *PeerCapability) {
	r.capMutex.Lock()
	defer r.capMutex.Unlock()
	
	capability.PeerID = peerID
	capability.LastSeen = time.Now()
	r.capabilities[peerID] = capability
	
	// Clear relevant cache entries
	r.invalidateCache()
}

// UnregisterPeer removes a peer from the registry
func (r *Resolver) UnregisterPeer(peerID peer.ID) {
	r.capMutex.Lock()
	defer r.capMutex.Unlock()
	
	delete(r.capabilities, peerID)
	
	// Clear relevant cache entries
	r.invalidateCache()
}

// UpdatePeerStatus updates a peer's status
func (r *Resolver) UpdatePeerStatus(peerID peer.ID, status string) {
	r.capMutex.Lock()
	defer r.capMutex.Unlock()
	
	if cap, exists := r.capabilities[peerID]; exists {
		cap.Status = status
		cap.LastSeen = time.Now()
	}
}

// Resolve resolves a bzzz:// URI to peer addresses
func (r *Resolver) Resolve(ctx context.Context, uri *BzzzURI, strategy ...ResolutionStrategy) (*ResolutionResult, error) {
	if uri == nil {
		return nil, fmt.Errorf("nil URI")
	}
	
	// Determine strategy
	resolveStrategy := r.defaultStrategy
	if len(strategy) > 0 {
		resolveStrategy = strategy[0]
	}
	
	// Check cache first
	cacheKey := r.getCacheKey(uri, resolveStrategy)
	if result := r.getFromCache(cacheKey); result != nil {
		return result, nil
	}
	
	// Perform resolution
	result, err := r.resolveURI(ctx, uri, resolveStrategy)
	if err != nil {
		return nil, err
	}
	
	// Cache result
	r.cacheResult(cacheKey, result)
	
	return result, nil
}

// ResolveString resolves a bzzz:// URI string to peer addresses  
func (r *Resolver) ResolveString(ctx context.Context, uriStr string, strategy ...ResolutionStrategy) (*ResolutionResult, error) {
	uri, err := ParseBzzzURI(uriStr)
	if err != nil {
		return nil, fmt.Errorf("failed to parse URI: %w", err)
	}
	
	return r.Resolve(ctx, uri, strategy...)
}

// resolveURI performs the actual URI resolution
func (r *Resolver) resolveURI(ctx context.Context, uri *BzzzURI, strategy ResolutionStrategy) (*ResolutionResult, error) {
	r.capMutex.RLock()
	defer r.capMutex.RUnlock()
	
	var matchingPeers []*PeerCapability
	
	// Find matching peers
	for _, cap := range r.capabilities {
		if r.peerMatches(cap, uri) {
			matchingPeers = append(matchingPeers, cap)
		}
	}
	
	if len(matchingPeers) == 0 {
		return &ResolutionResult{
			URI:           uri,
			Peers:         []*PeerAddress{},
			ResolvedAt:    time.Now(),
			ResolutionTTL: r.cacheTTL,
			Strategy:      string(strategy),
		}, nil
	}
	
	// Apply resolution strategy
	selectedPeers := r.applyStrategy(matchingPeers, strategy)
	
	// Convert to peer addresses
	var peerAddresses []*PeerAddress
	for i, cap := range selectedPeers {
		if i >= r.maxPeersPerResult {
			break
		}
		
		addr := &PeerAddress{
			PeerID:   cap.PeerID,
			Priority: r.calculatePriority(cap, uri),
			Metadata: map[string]interface{}{
				"agent":          cap.Agent,
				"role":           cap.Role,
				"specialization": cap.Specialization,
				"status":         cap.Status,
				"last_seen":      cap.LastSeen,
			},
		}
		
		// Get addresses from peerstore
		if r.peerstore != nil {
			addrs := r.peerstore.Addrs(cap.PeerID)
			for _, ma := range addrs {
				addr.Addresses = append(addr.Addresses, ma.String())
			}
		}
		
		peerAddresses = append(peerAddresses, addr)
	}
	
	return &ResolutionResult{
		URI:           uri,
		Peers:         peerAddresses,
		ResolvedAt:    time.Now(),
		ResolutionTTL: r.cacheTTL,
		Strategy:      string(strategy),
	}, nil
}

// peerMatches checks if a peer matches the URI criteria
func (r *Resolver) peerMatches(cap *PeerCapability, uri *BzzzURI) bool {
	// Check if peer is online
	if cap.Status == "offline" {
		return false
	}
	
	// Check agent match
	if !IsWildcard(uri.Agent) && !componentMatches(uri.Agent, cap.Agent) {
		return false
	}
	
	// Check role match
	if !IsWildcard(uri.Role) && !componentMatches(uri.Role, cap.Role) {
		return false
	}
	
	// Check project match (if specified in metadata)
	if !IsWildcard(uri.Project) {
		if project, exists := cap.Metadata["project"]; exists {
			if !componentMatches(uri.Project, project) {
				return false
			}
		}
	}
	
	// Check task capabilities (if peer has relevant capabilities)
	if !IsWildcard(uri.Task) {
		taskMatches := false
		for _, capability := range cap.Capabilities {
			if componentMatches(uri.Task, capability) {
				taskMatches = true
				break
			}
		}
		if !taskMatches {
			// Also check specialization
			if !componentMatches(uri.Task, cap.Specialization) {
				return false
			}
		}
	}
	
	return true
}

// applyStrategy applies the resolution strategy to matching peers
func (r *Resolver) applyStrategy(peers []*PeerCapability, strategy ResolutionStrategy) []*PeerCapability {
	switch strategy {
	case StrategyExact:
		// Return only exact matches (already filtered)
		return peers
		
	case StrategyPriority:
		// Sort by priority (calculated based on specificity and status)
		return r.sortByPriority(peers)
		
	case StrategyLoadBalance:
		// Sort by load (prefer less busy peers)
		return r.sortByLoad(peers)
		
	case StrategyBestMatch:
		fallthrough
	default:
		// Sort by best match score
		return r.sortByMatch(peers)
	}
}

// sortByPriority sorts peers by priority score
func (r *Resolver) sortByPriority(peers []*PeerCapability) []*PeerCapability {
	// Simple priority: online > working > busy, then by last seen
	result := make([]*PeerCapability, len(peers))
	copy(result, peers)
	
	// Sort by status priority and recency
	for i := 0; i < len(result)-1; i++ {
		for j := i + 1; j < len(result); j++ {
			iPriority := r.getStatusPriority(result[i].Status)
			jPriority := r.getStatusPriority(result[j].Status)
			
			if iPriority < jPriority || 
			   (iPriority == jPriority && result[i].LastSeen.Before(result[j].LastSeen)) {
				result[i], result[j] = result[j], result[i]
			}
		}
	}
	
	return result
}

// sortByLoad sorts peers by current load (prefer less busy)
func (r *Resolver) sortByLoad(peers []*PeerCapability) []*PeerCapability {
	result := make([]*PeerCapability, len(peers))
	copy(result, peers)
	
	// Sort by status (ready > working > busy)
	for i := 0; i < len(result)-1; i++ {
		for j := i + 1; j < len(result); j++ {
			iLoad := r.getLoadScore(result[i].Status)
			jLoad := r.getLoadScore(result[j].Status)
			
			if iLoad > jLoad {
				result[i], result[j] = result[j], result[i]
			}
		}
	}
	
	return result
}

// sortByMatch sorts peers by match quality
func (r *Resolver) sortByMatch(peers []*PeerCapability) []*PeerCapability {
	result := make([]*PeerCapability, len(peers))
	copy(result, peers)
	
	// Simple sorting - prefer online status and recent activity
	for i := 0; i < len(result)-1; i++ {
		for j := i + 1; j < len(result); j++ {
			if r.getMatchScore(result[i]) < r.getMatchScore(result[j]) {
				result[i], result[j] = result[j], result[i]
			}
		}
	}
	
	return result
}

// Helper functions for scoring
func (r *Resolver) getStatusPriority(status string) int {
	switch status {
	case "ready":
		return 3
	case "working":
		return 2
	case "busy":
		return 1
	default:
		return 0
	}
}

func (r *Resolver) getLoadScore(status string) int {
	switch status {
	case "ready":
		return 0 // Lowest load
	case "working":
		return 1
	case "busy":
		return 2 // Highest load
	default:
		return 3
	}
}

func (r *Resolver) getMatchScore(cap *PeerCapability) int {
	score := 0
	
	// Status contribution
	score += r.getStatusPriority(cap.Status) * 10
	
	// Recency contribution (more recent = higher score)
	timeSince := time.Since(cap.LastSeen)
	if timeSince < time.Minute {
		score += 5
	} else if timeSince < time.Hour {
		score += 3
	} else if timeSince < 24*time.Hour {
		score += 1
	}
	
	// Capability count contribution
	score += len(cap.Capabilities)
	
	return score
}

// calculatePriority calculates priority for a peer address
func (r *Resolver) calculatePriority(cap *PeerCapability, uri *BzzzURI) int {
	priority := 0
	
	// Exact matches get higher priority
	if cap.Agent == uri.Agent {
		priority += 4
	}
	if cap.Role == uri.Role {
		priority += 3
	}
	if cap.Specialization == uri.Task {
		priority += 2
	}
	
	// Status-based priority
	priority += r.getStatusPriority(cap.Status)
	
	return priority
}

// Cache management
func (r *Resolver) getCacheKey(uri *BzzzURI, strategy ResolutionStrategy) string {
	return fmt.Sprintf("%s:%s", uri.String(), strategy)
}

func (r *Resolver) getFromCache(key string) *ResolutionResult {
	r.cacheMutex.RLock()
	defer r.cacheMutex.RUnlock()
	
	if result, exists := r.cache[key]; exists {
		// Check if result is still valid
		if time.Since(result.ResolvedAt) < result.ResolutionTTL {
			return result
		}
		
		// Remove expired entry
		delete(r.cache, key)
	}
	
	return nil
}

func (r *Resolver) cacheResult(key string, result *ResolutionResult) {
	r.cacheMutex.Lock()
	defer r.cacheMutex.Unlock()
	
	r.cache[key] = result
}

func (r *Resolver) invalidateCache() {
	r.cacheMutex.Lock()
	defer r.cacheMutex.Unlock()
	
	// Clear entire cache on capability changes
	r.cache = make(map[string]*ResolutionResult)
}

// startCleanup starts background cache cleanup
func (r *Resolver) startCleanup() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	
	for range ticker.C {
		r.cleanupCache()
	}
}

func (r *Resolver) cleanupCache() {
	r.cacheMutex.Lock()
	defer r.cacheMutex.Unlock()
	
	now := time.Now()
	for key, result := range r.cache {
		if now.Sub(result.ResolvedAt) > result.ResolutionTTL {
			delete(r.cache, key)
		}
	}
}

// GetPeerCapabilities returns all registered peer capabilities
func (r *Resolver) GetPeerCapabilities() map[peer.ID]*PeerCapability {
	r.capMutex.RLock()
	defer r.capMutex.RUnlock()
	
	result := make(map[peer.ID]*PeerCapability)
	for id, cap := range r.capabilities {
		result[id] = cap
	}
	
	return result
}

// GetPeerCapability returns a specific peer's capabilities
func (r *Resolver) GetPeerCapability(peerID peer.ID) (*PeerCapability, bool) {
	r.capMutex.RLock()
	defer r.capMutex.RUnlock()
	
	cap, exists := r.capabilities[peerID]
	return cap, exists
}

// Close shuts down the resolver
func (r *Resolver) Close() error {
	// Clear all data
	r.capMutex.Lock()
	r.capabilities = make(map[peer.ID]*PeerCapability)
	r.capMutex.Unlock()
	
	r.cacheMutex.Lock()
	r.cache = make(map[string]*ResolutionResult)
	r.cacheMutex.Unlock()
	
	return nil
}