poimen-workflows/internal/clusters/cluster_manager.go

package clusters

import (
	"fmt"
	"sync"
	"time"
)

// ClusterInfo represents a Kubernetes cluster
type ClusterInfo struct {
	Name       string
	APIServer  string
	Healthy    bool
	LastCheck  time.Time
	Capacity   int // Max concurrent tasks
	Usage      int // Current task count
}

// ClusterManager manages multiple K8s clusters
type ClusterManager struct {
	mu       sync.RWMutex
	clusters map[string]*ClusterInfo
}

// NewClusterManager creates a new cluster manager
func NewClusterManager() *ClusterManager {
	return &ClusterManager{
		clusters: make(map[string]*ClusterInfo),
	}
}

// RegisterCluster registers a new cluster
func (cm *ClusterManager) RegisterCluster(name, apiServer string, capacity int) error {
	if name == "" || apiServer == "" {
		return fmt.Errorf("cluster name and API server required")
	}

	cm.mu.Lock()
	defer cm.mu.Unlock()

	if _, exists := cm.clusters[name]; exists {
		return fmt.Errorf("cluster already registered: %s", name)
	}

	cm.clusters[name] = &ClusterInfo{
		Name:      name,
		APIServer: apiServer,
		Healthy:   true,
		LastCheck: time.Now(),
		Capacity:  capacity,
		Usage:     0,
	}

	return nil
}

// UnregisterCluster removes a cluster
func (cm *ClusterManager) UnregisterCluster(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	if _, exists := cm.clusters[name]; !exists {
		return fmt.Errorf("cluster not found: %s", name)
	}

	delete(cm.clusters, name)
	return nil
}

// GetCluster retrieves cluster info
func (cm *ClusterManager) GetCluster(name string) (*ClusterInfo, bool) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	cluster, exists := cm.clusters[name]
	return cluster, exists
}

// ListClusters returns all registered clusters
func (cm *ClusterManager) ListClusters() map[string]*ClusterInfo {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	result := make(map[string]*ClusterInfo)
	for name, cluster := range cm.clusters {
		result[name] = cluster
	}

	return result
}

// HealthCheck checks cluster health
func (cm *ClusterManager) HealthCheck(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	cluster, exists := cm.clusters[name]
	if !exists {
		return fmt.Errorf("cluster not found: %s", name)
	}

	// Simulate health check (in production, query API server)
	cluster.Healthy = true
	cluster.LastCheck = time.Now()

	return nil
}

// MarkUnhealthy marks a cluster as unhealthy
func (cm *ClusterManager) MarkUnhealthy(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	cluster, exists := cm.clusters[name]
	if !exists {
		return fmt.Errorf("cluster not found: %s", name)
	}

	cluster.Healthy = false
	return nil
}

// AllocateTask allocates a task to a cluster
func (cm *ClusterManager) AllocateTask(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	cluster, exists := cm.clusters[name]
	if !exists {
		return fmt.Errorf("cluster not found: %s", name)
	}

	if !cluster.Healthy {
		return fmt.Errorf("cluster not healthy: %s", name)
	}

	if cluster.Usage >= cluster.Capacity {
		return fmt.Errorf("cluster at capacity: %s", name)
	}

	cluster.Usage++
	return nil
}

// ReleaseTask releases a task from a cluster
func (cm *ClusterManager) ReleaseTask(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	cluster, exists := cm.clusters[name]
	if !exists {
		return fmt.Errorf("cluster not found: %s", name)
	}

	if cluster.Usage > 0 {
		cluster.Usage--
	}

	return nil
}

// FindBestCluster finds the cluster with most available capacity
func (cm *ClusterManager) FindBestCluster() (string, error) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	var bestCluster string
	maxCapacity := 0

	for name, cluster := range cm.clusters {
		if cluster.Healthy {
			available := cluster.Capacity - cluster.Usage
			if available > maxCapacity {
				bestCluster = name
				maxCapacity = available
			}
		}
	}

	if bestCluster == "" {
		return "", fmt.Errorf("no healthy clusters available")
	}

	return bestCluster, nil
}

// GetCapacitySummary returns capacity summary
func (cm *ClusterManager) GetCapacitySummary() map[string]interface{} {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	totalCapacity := 0
	totalUsage := 0
	healthyCount := 0

	for _, cluster := range cm.clusters {
		totalCapacity += cluster.Capacity
		totalUsage += cluster.Usage
		if cluster.Healthy {
			healthyCount++
		}
	}

	return map[string]interface{}{
		"total_capacity": totalCapacity,
		"total_usage":    totalUsage,
		"healthy_clusters": healthyCount,
		"total_clusters": len(cm.clusters),
	}
}

// GetHealthStatus returns health status for all clusters
func (cm *ClusterManager) GetHealthStatus() map[string]bool {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	result := make(map[string]bool)
	for name, cluster := range cm.clusters {
		result[name] = cluster.Healthy
	}

	return result
}