Control-Center/go-backend/internal/store/agent.go

// Package store provides thread-safe in-memory data stores for the
// Control Center API. These will be replaced with PostgreSQL-backed
// implementations once CUB-120 (schema design) is complete.
package store

import (
	"sort"
	"sync"
	"time"

	"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/models"
	"github.com/google/uuid"
)

// AgentStore provides thread-safe CRUD operations for agents.
type AgentStore struct {
	mu      sync.RWMutex
	agents  map[string]models.AgentCardData
	history map[string][]models.AgentStatusHistoryEntry // agentID -> history
}

// NewAgentStore returns an initialized AgentStore.
func NewAgentStore() *AgentStore {
	return &AgentStore{
		agents:  make(map[string]models.AgentCardData),
		history: make(map[string][]models.AgentStatusHistoryEntry),
	}
}

// List returns all agents, optionally filtered by status.
func (s *AgentStore) List(statusFilter models.AgentStatus) []models.AgentCardData {
	s.mu.RLock()
	defer s.mu.RUnlock()

	result := make([]models.AgentCardData, 0, len(s.agents))
	for _, a := range s.agents {
		if statusFilter != "" && a.Status != statusFilter {
			continue
		}
		result = append(result, a)
	}

	// Sort by display name for consistent output.
	sort.Slice(result, func(i, j int) bool {
		return result[i].DisplayName < result[j].DisplayName
	})
	return result
}

// Get returns a single agent by ID, or false if not found.
func (s *AgentStore) Get(id string) (models.AgentCardData, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	a, ok := s.agents[id]
	return a, ok
}

// Create inserts a new agent. Returns false if the ID already exists.
func (s *AgentStore) Create(a models.AgentCardData) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.agents[a.ID]; exists {
		return false
	}
	if a.LastActivity == "" {
		a.LastActivity = time.Now().UTC().Format(time.RFC3339)
	}
	s.agents[a.ID] = a

	// Record initial history entry.
	s.appendHistoryLocked(a.ID, models.AgentStatusHistoryEntry{
		ID:        uuid.New().String(),
		AgentID:   a.ID,
		Status:    a.Status,
		Task:      a.CurrentTask,
		Timestamp: a.LastActivity,
	})
	return true
}

// Update applies partial updates to an agent. Returns the updated agent or false if not found.
func (s *AgentStore) Update(id string, req models.UpdateAgentRequest) (models.AgentCardData, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	a, ok := s.agents[id]
	if !ok {
		return models.AgentCardData{}, false
	}

	prevStatus := a.Status
	prevTask := a.CurrentTask

	if req.Status != nil {
		a.Status = *req.Status
	}
	if req.CurrentTask != nil {
		a.CurrentTask = req.CurrentTask
	}
	if req.TaskProgress != nil {
		a.TaskProgress = req.TaskProgress
	}
	if req.TaskElapsed != nil {
		a.TaskElapsed = req.TaskElapsed
	}
	if req.Channel != nil {
		a.Channel = *req.Channel
	}
	if req.ErrorMessage != nil {
		a.ErrorMessage = req.ErrorMessage
	}
	a.LastActivity = time.Now().UTC().Format(time.RFC3339)
	s.agents[id] = a

	// Record history entry if status or task changed.
	if (req.Status != nil && *req.Status != prevStatus) || (req.CurrentTask != nil && prevTask == nil) ||
		(req.CurrentTask != nil && prevTask != nil && *req.CurrentTask != *prevTask) {
		status := a.Status
		if req.Status != nil {
			status = *req.Status
		}
		s.appendHistoryLocked(id, models.AgentStatusHistoryEntry{
			ID:        uuid.New().String(),
			AgentID:   id,
			Status:    status,
			Task:      a.CurrentTask,
			Timestamp: a.LastActivity,
		})
	}

	return a, true
}

// Delete removes an agent. Returns true if the agent existed.
func (s *AgentStore) Delete(id string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.agents[id]; !ok {
		return false
	}
	delete(s.agents, id)
	delete(s.history, id)
	return true
}

// History returns the status history for an agent, newest first.
func (s *AgentStore) History(agentID string) []models.AgentStatusHistoryEntry {
	s.mu.RLock()
	defer s.mu.RUnlock()

	entries, ok := s.history[agentID]
	if !ok {
		return nil
	}
	// Return a defensive copy, sorted newest first (by index when timestamps tie).
	result := make([]models.AgentStatusHistoryEntry, len(entries))
	copy(result, entries)
	sort.SliceStable(result, func(i, j int) bool {
		if result[i].Timestamp == result[j].Timestamp {
			return i > j // later index = newer when timestamps match
		}
		return result[i].Timestamp > result[j].Timestamp
	})
	return result
}

// Count returns the total number of agents.
func (s *AgentStore) Count() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.agents)
}

// appendHistoryLocked adds a history entry. Caller must hold s.mu (write lock).
func (s *AgentStore) appendHistoryLocked(agentID string, entry models.AgentStatusHistoryEntry) {
	s.history[agentID] = append(s.history[agentID], entry)
}