HellreigN/backend/internal/graph/graph.go

package graph

import (
	"fmt"
	"sort"
)

// DepCondition represents how a service waits for a dependency.
type DepCondition string

const (
	Started               DepCondition = "started"
	Healthy               DepCondition = "healthy"
	CompletedSuccessfully DepCondition = "completed_successfully"
)

// ServiceRef uniquely identifies a service across nodes.
// If NodeID is empty, it refers to a service in the same node.
type ServiceRef struct {
	NodeID string `json:"node_id,omitempty"`
	Name   string `json:"name"`
}

// String returns a human-readable reference like "node:service" or just "service".
func (r ServiceRef) String() string {
	if r.NodeID != "" {
		return r.NodeID + ":" + r.Name
	}
	return r.Name
}

// Dependency declares that a service depends on another service (possibly in a different node).
type Dependency struct {
	Target    ServiceRef   `json:"target"`
	Condition DepCondition `json:"condition"`
}

// Service represents a named service within a node with its dependency declarations.
type Service struct {
	Name         string       `json:"name"`
	Dependencies []Dependency `json:"dependencies,omitempty"`
}

// Node represents a logical grouping of services (e.g., a server or cluster).
type Node struct {
	ID       string    `json:"id"`
	Services []*Service `json:"services"`
}

// Graph holds nodes, services, and computes dependency order.
type Graph struct {
	nodes map[string]*Node
	// adj[key] = list of services that key depends on
	// key format: "nodeID:serviceName"
	adj map[string][]ServiceRef
}

func New() *Graph {
	return &Graph{
		nodes: make(map[string]*Node),
		adj:   make(map[string][]ServiceRef),
	}
}

// AddNode adds a node to the graph.
func (g *Graph) AddNode(nodeID string) *Node {
	if n, ok := g.nodes[nodeID]; ok {
		return n
	}
	n := &Node{ID: nodeID}
	g.nodes[nodeID] = n
	return n
}

// AddService adds a service to a node.
func (g *Graph) AddService(nodeID string, svc *Service) {
	node := g.AddNode(nodeID)
	node.Services = append(node.Services, svc)
	key := nodeID + ":" + svc.Name
	g.adj[key] = nil
}

// ResolveRef resolves a ServiceRef to its full "nodeID:serviceName" key.
// If ref.NodeID is empty, it's resolved relative to the given sourceNodeID.
func (g *Graph) ResolveRef(ref ServiceRef, sourceNodeID string) (string, error) {
	nodeID := ref.NodeID
	if nodeID == "" {
		nodeID = sourceNodeID
	}
	key := nodeID + ":" + ref.Name
	if _, ok := g.adj[key]; !ok {
		return "", fmt.Errorf("unknown service %q", key)
	}
	return key, nil
}

// AddDependency adds a dependency: source service depends on target service.
func (g *Graph) AddDependency(sourceNodeID, sourceName string, dep Dependency) error {
	srcKey := sourceNodeID + ":" + sourceName
	if _, ok := g.adj[srcKey]; !ok {
		return fmt.Errorf("unknown source service %q", srcKey)
	}

	if _, err := g.ResolveRef(dep.Target, sourceNodeID); err != nil {
		return fmt.Errorf("dependency target invalid: %w", err)
	}

	g.adj[srcKey] = append(g.adj[srcKey], dep.Target)

	// Also update the Service struct for serialization
	node, ok := g.nodes[sourceNodeID]
	if !ok {
		return nil
	}
	for _, svc := range node.Services {
		if svc.Name == sourceName {
			svc.Dependencies = append(svc.Dependencies, dep)
			break
		}
	}
	return nil
}

// HasCycle detects if the dependency graph contains a cycle.
func (g *Graph) HasCycle() bool {
	const (
		white = 0
		gray  = 1
		black = 2
	)
	color := make(map[string]int)
	for key := range g.adj {
		color[key] = white
	}

	var dfs func(string) bool
	dfs = func(u string) bool {
		color[u] = gray
		for _, depRef := range g.adj[u] {
			v, _ := g.ResolveRef(depRef, nodeIDFromKey(u))
			if color[v] == gray {
				return true
			}
			if color[v] == white && dfs(v) {
				return true
			}
		}
		color[u] = black
		return false
	}

	for key := range g.adj {
		if color[key] == white {
			if dfs(key) {
				return true
			}
		}
	}
	return false
}

// TopologicalSort returns services in startup order (dependencies first).
// Returns a flat list of "nodeID:serviceName" keys.
func (g *Graph) TopologicalSort() ([]string, error) {
	if g.HasCycle() {
		return nil, fmt.Errorf("dependency cycle detected")
	}

	var result []string
	visited := make(map[string]bool)

	var dfs func(string)
	dfs = func(u string) {
		if visited[u] {
			return
		}
		visited[u] = true
		for _, depRef := range g.adj[u] {
			v, _ := g.ResolveRef(depRef, nodeIDFromKey(u))
			dfs(v)
		}
		result = append(result, u)
	}

	keys := make([]string, 0, len(g.adj))
	for k := range g.adj {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		dfs(k)
	}

	return result, nil
}

// GetNode returns a node by ID.
func (g *Graph) GetNode(id string) (*Node, bool) {
	n, ok := g.nodes[id]
	return n, ok
}

// GetService returns a service by node ID and name.
func (g *Graph) GetService(nodeID, name string) (*Service, bool) {
	node, ok := g.nodes[nodeID]
	if !ok {
		return nil, false
	}
	for _, s := range node.Services {
		if s.Name == name {
			return s, true
		}
	}
	return nil, false
}

// Nodes returns all nodes sorted by ID.
func (g *Graph) Nodes() []*Node {
	result := make([]*Node, 0, len(g.nodes))
	for _, n := range g.nodes {
		result = append(result, n)
	}
	sort.Slice(result, func(i, j int) bool {
		return result[i].ID < result[j].ID
	})
	return result
}

// nodeIDFromKey extracts the node ID from a "nodeID:serviceName" key.
func nodeIDFromKey(key string) string {
	for i := 0; i < len(key); i++ {
		if key[i] == ':' {
			return key[:i]
		}
	}
	return ""
}