WHOOSH/internal/agents/registry.go

package agents

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

	"github.com/chorus-services/whoosh/internal/p2p"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/rs/zerolog/log"
)

// Registry manages agent registration and synchronization with the database
type Registry struct {
	db        *pgxpool.Pool
	discovery *p2p.Discovery
	stopCh    chan struct{}
	ctx       context.Context
	cancel    context.CancelFunc
}

// NewRegistry creates a new agent registry service
func NewRegistry(db *pgxpool.Pool, discovery *p2p.Discovery) *Registry {
	ctx, cancel := context.WithCancel(context.Background())
	
	return &Registry{
		db:        db,
		discovery: discovery,
		stopCh:    make(chan struct{}),
		ctx:       ctx,
		cancel:    cancel,
	}
}

// Start begins the agent registry synchronization
func (r *Registry) Start() error {
	log.Info().Msg("🔄 Starting CHORUS agent registry synchronization")
	
	// Start periodic synchronization of discovered agents with database
	go r.syncDiscoveredAgents()
	
	return nil
}

// Stop shuts down the agent registry
func (r *Registry) Stop() error {
	log.Info().Msg("🔄 Stopping CHORUS agent registry synchronization")
	
	r.cancel()
	close(r.stopCh)
	
	return nil
}

// syncDiscoveredAgents periodically syncs P2P discovered agents to database
func (r *Registry) syncDiscoveredAgents() {
	// Initial sync
	r.performSync()
	
	// Then sync every 30 seconds
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	
	for {
		select {
		case <-r.ctx.Done():
			return
		case <-ticker.C:
			r.performSync()
		}
	}
}

// performSync synchronizes discovered agents with the database
func (r *Registry) performSync() {
	discoveredAgents := r.discovery.GetAgents()
	
	log.Debug().
		Int("discovered_count", len(discoveredAgents)).
		Msg("Synchronizing discovered agents with database")
	
	for _, agent := range discoveredAgents {
		err := r.upsertAgent(r.ctx, agent)
		if err != nil {
			log.Error().
				Err(err).
				Str("agent_id", agent.ID).
				Msg("Failed to sync agent to database")
		}
	}
	
	// Clean up agents that are no longer discovered
	err := r.markOfflineAgents(r.ctx, discoveredAgents)
	if err != nil {
		log.Error().
			Err(err).
			Msg("Failed to mark offline agents")
	}
}

// upsertAgent inserts or updates an agent in the database
func (r *Registry) upsertAgent(ctx context.Context, agent *p2p.Agent) error {
	// Convert capabilities to JSON
	capabilitiesJSON, err := json.Marshal(agent.Capabilities)
	if err != nil {
		return fmt.Errorf("failed to marshal capabilities: %w", err)
	}
	
	// Create performance metrics
	performanceMetrics := map[string]interface{}{
		"tasks_completed": agent.TasksCompleted,
		"current_team":    agent.CurrentTeam,
		"model":          agent.Model,
		"cluster_id":     agent.ClusterID,
		"p2p_addr":       agent.P2PAddr,
	}
	metricsJSON, err := json.Marshal(performanceMetrics)
	if err != nil {
		return fmt.Errorf("failed to marshal performance metrics: %w", err)
	}
	
	// Map P2P status to database status
	dbStatus := r.mapStatusToDatabase(agent.Status)
	
	// Use upsert query to insert or update
	query := `
		INSERT INTO agents (id, name, endpoint_url, capabilities, status, last_seen, performance_metrics, current_tasks, success_rate)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
		ON CONFLICT (id) 
		DO UPDATE SET
			name = EXCLUDED.name,
			endpoint_url = EXCLUDED.endpoint_url,
			capabilities = EXCLUDED.capabilities,
			status = EXCLUDED.status,
			last_seen = EXCLUDED.last_seen,
			performance_metrics = EXCLUDED.performance_metrics,
			current_tasks = EXCLUDED.current_tasks,
			updated_at = NOW()
		RETURNING id
	`
	
	// Generate UUID from agent ID for database consistency
	agentUUID, err := r.generateConsistentUUID(agent.ID)
	if err != nil {
		return fmt.Errorf("failed to generate UUID: %w", err)
	}
	
	var resultID uuid.UUID
	err = r.db.QueryRow(ctx, query,
		agentUUID,           // id
		agent.Name,          // name
		agent.Endpoint,      // endpoint_url
		capabilitiesJSON,    // capabilities
		dbStatus,            // status
		agent.LastSeen,      // last_seen
		metricsJSON,         // performance_metrics
		r.getCurrentTaskCount(agent), // current_tasks
		r.calculateSuccessRate(agent), // success_rate
	).Scan(&resultID)
	
	if err != nil {
		return fmt.Errorf("failed to upsert agent: %w", err)
	}
	
	log.Debug().
		Str("agent_id", agent.ID).
		Str("db_uuid", resultID.String()).
		Str("status", dbStatus).
		Msg("Synced agent to database")
	
	return nil
}

// markOfflineAgents marks agents as offline if they're no longer discovered
func (r *Registry) markOfflineAgents(ctx context.Context, discoveredAgents []*p2p.Agent) error {
	// Build list of currently discovered agent IDs
	discoveredIDs := make([]string, len(discoveredAgents))
	for i, agent := range discoveredAgents {
		discoveredIDs[i] = agent.ID
	}
	
	// Convert to UUIDs for database query
	discoveredUUIDs := make([]uuid.UUID, len(discoveredIDs))
	for i, id := range discoveredIDs {
		uuid, err := r.generateConsistentUUID(id)
		if err != nil {
			return fmt.Errorf("failed to generate UUID for %s: %w", id, err)
		}
		discoveredUUIDs[i] = uuid
	}
	
	// If no agents discovered, don't mark all as offline (could be temporary network issue)
	if len(discoveredUUIDs) == 0 {
		return nil
	}
	
	// Mark agents as offline if they haven't been seen and aren't in discovered list
	query := `
		UPDATE agents 
		SET status = 'offline', updated_at = NOW()
		WHERE status != 'offline' 
		  AND last_seen < NOW() - INTERVAL '2 minutes'
		  AND id != ALL($1)
	`
	
	result, err := r.db.Exec(ctx, query, discoveredUUIDs)
	if err != nil {
		return fmt.Errorf("failed to mark offline agents: %w", err)
	}
	
	rowsAffected := result.RowsAffected()
	if rowsAffected > 0 {
		log.Info().
			Int64("agents_marked_offline", rowsAffected).
			Msg("Marked agents as offline")
	}
	
	return nil
}

// mapStatusToDatabase maps P2P status to database status values
func (r *Registry) mapStatusToDatabase(p2pStatus string) string {
	switch p2pStatus {
	case "online":
		return "available"
	case "idle":
		return "idle"
	case "working":
		return "busy"
	default:
		return "available"
	}
}

// getCurrentTaskCount estimates current task count based on status
func (r *Registry) getCurrentTaskCount(agent *p2p.Agent) int {
	switch agent.Status {
	case "working":
		return 1
	case "idle", "online":
		return 0
	default:
		return 0
	}
}

// calculateSuccessRate calculates success rate based on tasks completed
func (r *Registry) calculateSuccessRate(agent *p2p.Agent) float64 {
	// For MVP, assume high success rate for all agents
	// In production, this would be calculated from actual task outcomes
	if agent.TasksCompleted > 0 {
		return 0.85 + (float64(agent.TasksCompleted)*0.01) // Success rate increases with experience
	}
	return 0.75 // Default for new agents
}

// generateConsistentUUID generates a consistent UUID from a string ID
// This ensures the same agent ID always maps to the same UUID
func (r *Registry) generateConsistentUUID(agentID string) (uuid.UUID, error) {
	// Use UUID v5 (name-based) to generate consistent UUIDs
	// This ensures the same agent ID always produces the same UUID
	namespace := uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8") // DNS namespace UUID
	return uuid.NewSHA1(namespace, []byte(agentID)), nil
}

// GetAvailableAgents returns agents that are available for task assignment
func (r *Registry) GetAvailableAgents(ctx context.Context) ([]*DatabaseAgent, error) {
	query := `
		SELECT id, name, endpoint_url, capabilities, status, last_seen, 
		       performance_metrics, current_tasks, success_rate, created_at, updated_at
		FROM agents 
		WHERE status IN ('available', 'idle') 
		  AND last_seen > NOW() - INTERVAL '5 minutes'
		ORDER BY success_rate DESC, current_tasks ASC
	`
	
	rows, err := r.db.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to query available agents: %w", err)
	}
	defer rows.Close()
	
	var agents []*DatabaseAgent
	
	for rows.Next() {
		agent := &DatabaseAgent{}
		var capabilitiesJSON, metricsJSON []byte
		
		err := rows.Scan(
			&agent.ID, &agent.Name, &agent.EndpointURL, &capabilitiesJSON,
			&agent.Status, &agent.LastSeen, &metricsJSON,
			&agent.CurrentTasks, &agent.SuccessRate,
			&agent.CreatedAt, &agent.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan agent row: %w", err)
		}
		
		// Parse JSON fields
		if len(capabilitiesJSON) > 0 {
			json.Unmarshal(capabilitiesJSON, &agent.Capabilities)
		}
		if len(metricsJSON) > 0 {
			json.Unmarshal(metricsJSON, &agent.PerformanceMetrics)
		}
		
		agents = append(agents, agent)
	}
	
	return agents, rows.Err()
}

// DatabaseAgent represents an agent as stored in the database
type DatabaseAgent struct {
	ID                 uuid.UUID              `json:"id" db:"id"`
	Name               string                 `json:"name" db:"name"`
	EndpointURL        string                 `json:"endpoint_url" db:"endpoint_url"`
	Capabilities       map[string]interface{} `json:"capabilities" db:"capabilities"`
	Status             string                 `json:"status" db:"status"`
	LastSeen           time.Time              `json:"last_seen" db:"last_seen"`
	PerformanceMetrics map[string]interface{} `json:"performance_metrics" db:"performance_metrics"`
	CurrentTasks       int                    `json:"current_tasks" db:"current_tasks"`
	SuccessRate        float64                `json:"success_rate" db:"success_rate"`
	CreatedAt          time.Time              `json:"created_at" db:"created_at"`
	UpdatedAt          time.Time              `json:"updated_at" db:"updated_at"`
}