CHORUS/pkg/protocol/integration.go

package protocol

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"chorus/pkg/config"
	"chorus/pkg/dht"
	"chorus/p2p"
	"github.com/libp2p/go-libp2p/core/peer"
)

// ProtocolManager manages the CHORUS v2 protocol components
type ProtocolManager struct {
	config   *config.Config
	node     *p2p.Node
	resolver *Resolver
	enabled  bool
	
	// Local peer information
	localPeer *PeerCapability
}

// NewProtocolManager creates a new protocol manager
func NewProtocolManager(cfg *config.Config, node *p2p.Node) (*ProtocolManager, error) {
	if cfg == nil || node == nil {
		return nil, fmt.Errorf("config and node are required")
	}
	
	pm := &ProtocolManager{
		config:  cfg,
		node:    node,
		enabled: cfg.V2.Enabled,
	}
	
	// Only initialize if v2 protocol is enabled
	if pm.enabled {
		if err := pm.initialize(); err != nil {
			return nil, fmt.Errorf("failed to initialize protocol manager: %w", err)
		}
	}
	
	return pm, nil
}

// initialize sets up the protocol components
func (pm *ProtocolManager) initialize() error {
	// Create resolver
	resolverOpts := []ResolverOption{
		WithCacheTTL(pm.config.V2.URIResolution.CacheTTL),
		WithMaxPeersPerResult(pm.config.V2.URIResolution.MaxPeersPerResult),
	}
	
	// Set default strategy
	switch pm.config.V2.URIResolution.DefaultStrategy {
	case "exact":
		resolverOpts = append(resolverOpts, WithDefaultStrategy(StrategyExact))
	case "priority":
		resolverOpts = append(resolverOpts, WithDefaultStrategy(StrategyPriority))
	case "load_balance":
		resolverOpts = append(resolverOpts, WithDefaultStrategy(StrategyLoadBalance))
	default:
		resolverOpts = append(resolverOpts, WithDefaultStrategy(StrategyBestMatch))
	}
	
	pm.resolver = NewResolver(pm.node.Host().Peerstore(), resolverOpts...)
	
	// Initialize local peer information
	pm.localPeer = &PeerCapability{
		PeerID:         pm.node.ID(),
		Agent:          pm.config.Agent.ID,
		Role:           pm.config.Agent.Role,
		Capabilities:   pm.config.Agent.Capabilities,
		Models:         pm.config.Agent.Models,
		Specialization: pm.config.Agent.Specialization,
		LastSeen:       time.Now(),
		Status:         "ready",
		Metadata:       make(map[string]string),
	}
	
	// Add project information if available
	if project := pm.getProjectFromConfig(); project != "" {
		pm.localPeer.Metadata["project"] = project
	}
	
	// Register local peer
	pm.resolver.RegisterPeer(pm.node.ID(), pm.localPeer)
	
	return nil
}

// IsEnabled returns whether the v2 protocol is enabled
func (pm *ProtocolManager) IsEnabled() bool {
	return pm.enabled
}

// ResolveURI resolves a CHORUS:// URI to peer addresses
func (pm *ProtocolManager) ResolveURI(ctx context.Context, uriStr string) (*ResolutionResult, error) {
	if !pm.enabled {
		return nil, fmt.Errorf("v2 protocol not enabled")
	}
	
	return pm.resolver.ResolveString(ctx, uriStr)
}

// RegisterPeer registers a peer's capabilities
func (pm *ProtocolManager) RegisterPeer(peerID peer.ID, capabilities *PeerCapability) {
	if !pm.enabled {
		return
	}
	
	pm.resolver.RegisterPeer(peerID, capabilities)
	
	// Announce to DHT if enabled
	if pm.node.IsDHTEnabled() {
		pm.announcePeerToDHT(context.Background(), capabilities)
	}
}

// UpdateLocalPeerStatus updates the local peer's status
func (pm *ProtocolManager) UpdateLocalPeerStatus(status string) {
	if !pm.enabled {
		return
	}
	
	pm.localPeer.Status = status
	pm.localPeer.LastSeen = time.Now()
	
	pm.resolver.RegisterPeer(pm.node.ID(), pm.localPeer)
}

// GetLocalPeer returns the local peer information
func (pm *ProtocolManager) GetLocalPeer() *PeerCapability {
	return pm.localPeer
}

// GetAllPeers returns all known peers
func (pm *ProtocolManager) GetAllPeers() map[peer.ID]*PeerCapability {
	if !pm.enabled {
		return make(map[peer.ID]*PeerCapability)
	}
	
	return pm.resolver.GetPeerCapabilities()
}

// HandlePeerCapabilityMessage handles incoming peer capability messages
func (pm *ProtocolManager) HandlePeerCapabilityMessage(peerID peer.ID, data []byte) error {
	if !pm.enabled {
		return nil // Silently ignore if v2 not enabled
	}
	
	var capability PeerCapability
	if err := json.Unmarshal(data, &capability); err != nil {
		return fmt.Errorf("failed to unmarshal capability message: %w", err)
	}
	
	capability.PeerID = peerID
	capability.LastSeen = time.Now()
	
	pm.resolver.RegisterPeer(peerID, &capability)
	
	return nil
}

// AnnounceCapabilities announces the local peer's capabilities
func (pm *ProtocolManager) AnnounceCapabilities() error {
	if !pm.enabled {
		return nil
	}
	
	// Update local peer information
	pm.localPeer.LastSeen = time.Now()
	
	// Announce to DHT if enabled
	if pm.node.IsDHTEnabled() {
		return pm.announcePeerToDHT(context.Background(), pm.localPeer)
	}
	
	return nil
}

// announcePeerToDHT announces a peer's capabilities to the DHT
func (pm *ProtocolManager) announcePeerToDHT(ctx context.Context, capability *PeerCapability) error {
	dht := pm.node.DHT()
	if dht == nil {
		return fmt.Errorf("DHT not available")
	}
	
	// Register peer with role-based and capability-based keys
	if capability.Role != "" {
		dht.RegisterPeer(capability.PeerID, capability.Agent, capability.Role, capability.Capabilities)
		if err := dht.AnnounceRole(ctx, capability.Role); err != nil {
			// Log error but don't fail
		}
	}
	
	// Announce each capability
	for _, cap := range capability.Capabilities {
		if err := dht.AnnounceCapability(ctx, cap); err != nil {
			// Log error but don't fail
		}
	}
	
	// Announce general peer presence
	if err := dht.Provide(ctx, "CHORUS:peer"); err != nil {
		// Log error but don't fail
	}
	
	return nil
}

// FindPeersByRole finds peers with a specific role
func (pm *ProtocolManager) FindPeersByRole(ctx context.Context, role string) ([]*PeerCapability, error) {
	if !pm.enabled {
		return nil, fmt.Errorf("v2 protocol not enabled")
	}
	
	// First try DHT if available
	if pm.node.IsDHTEnabled() {
		dhtPeers, err := pm.node.DHT().FindPeersByRole(ctx, role)
		if err == nil && len(dhtPeers) > 0 {
			// Convert DHT peer info to capabilities
			var capabilities []*PeerCapability
			for _, dhtPeer := range dhtPeers {
				cap := &PeerCapability{
					PeerID:   dhtPeer.ID,
					Agent:    dhtPeer.Agent,
					Role:     dhtPeer.Role,
					LastSeen: dhtPeer.LastSeen,
					Metadata: make(map[string]string),
				}
				capabilities = append(capabilities, cap)
			}
			return capabilities, nil
		}
	}
	
	// Fall back to local resolver
	var result []*PeerCapability
	for _, peer := range pm.resolver.GetPeerCapabilities() {
		if peer.Role == role || role == "*" {
			result = append(result, peer)
		}
	}
	
	return result, nil
}

// ValidateURI validates a CHORUS:// URI
func (pm *ProtocolManager) ValidateURI(uriStr string) error {
	if !pm.enabled {
		return fmt.Errorf("v2 protocol not enabled")
	}
	
	_, err := ParseBzzzURI(uriStr)
	return err
}

// CreateURI creates a CHORUS:// URI with the given components
func (pm *ProtocolManager) CreateURI(agent, role, project, task, path string) (*BzzzURI, error) {
	if !pm.enabled {
		return nil, fmt.Errorf("v2 protocol not enabled")
	}
	
	// Use configured defaults if components are empty
	if agent == "" {
		agent = pm.config.V2.SemanticAddressing.DefaultAgent
	}
	if role == "" {
		role = pm.config.V2.SemanticAddressing.DefaultRole
	}
	if project == "" {
		project = pm.config.V2.SemanticAddressing.DefaultProject
	}
	
	return NewBzzzURI(agent, role, project, task, path), nil
}

// GetFeatureFlags returns the current feature flags
func (pm *ProtocolManager) GetFeatureFlags() map[string]bool {
	return pm.config.V2.FeatureFlags
}

// IsFeatureEnabled checks if a specific feature is enabled
func (pm *ProtocolManager) IsFeatureEnabled(feature string) bool {
	if !pm.enabled {
		return false
	}
	
	enabled, exists := pm.config.V2.FeatureFlags[feature]
	return exists && enabled
}

// Close shuts down the protocol manager
func (pm *ProtocolManager) Close() error {
	if pm.resolver != nil {
		return pm.resolver.Close()
	}
	return nil
}

// getProjectFromConfig extracts project information from configuration
func (pm *ProtocolManager) getProjectFromConfig() string {
	// Try to infer project from agent ID or other configuration
	if pm.config.Agent.ID != "" {
		parts := strings.Split(pm.config.Agent.ID, "-")
		if len(parts) > 0 {
			return parts[0]
		}
	}
	
	// Default project if none can be inferred
	return "CHORUS"
}

// GetStats returns protocol statistics
func (pm *ProtocolManager) GetStats() map[string]interface{} {
	stats := map[string]interface{}{
		"enabled":     pm.enabled,
		"local_peer":  pm.localPeer,
		"known_peers": len(pm.resolver.GetPeerCapabilities()),
	}
	
	if pm.node.IsDHTEnabled() {
		dht := pm.node.DHT()
		stats["dht_enabled"] = true
		stats["dht_bootstrapped"] = dht.IsBootstrapped()
		stats["dht_size"] = dht.GetDHTSize()
		stats["dht_connected_peers"] = len(dht.GetConnectedPeers())
	} else {
		stats["dht_enabled"] = false
	}
	
	return stats
}