kubeviz/internal/graph/graph.go

package graph

import (
	"fmt"
	"sort"
	"strings"

	"kubeviz/internal/model"
	"kubeviz/internal/parser"
)

type Filters struct {
	Namespace       string
	Kind            string
	Query           string
	FocusID         string
	Relations       map[string]bool
	GroupBy         string
	CollapsedGroups map[string]bool
}

func BuildGraph(dataset *model.Dataset, filters Filters, healthHints map[string]string) model.GraphResponse {
	if dataset == nil {
		return model.GraphResponse{Stats: model.GraphStats{Kinds: map[string]int{}}}
	}

	nodes := make([]model.GraphNode, 0, len(dataset.Resources))
	edges := make([]model.GraphEdge, 0)
	kindStats := map[string]int{}

	for _, r := range dataset.Resources {
		health := "unknown"
		if h, ok := healthHints[r.ID]; ok {
			health = h
		}
		node := model.GraphNode{
			ID:          r.ID,
			Kind:        r.Kind,
			Name:        r.Name,
			Namespace:   r.Namespace,
			Labels:      r.Labels,
			HealthHint:  health,
			IsSensitive: r.IsSensitive,
		}
		if !matchesResourceFilters(node, filters) {
			continue
		}
		nodes = append(nodes, node)
		kindStats[node.Kind]++
	}

	nodeSet := make(map[string]struct{}, len(nodes))
	for _, n := range nodes {
		nodeSet[n.ID] = struct{}{}
	}

	for _, r := range dataset.Resources {
		if _, ok := nodeSet[r.ID]; !ok {
			continue
		}
		for _, ref := range r.References {
			target := resolveReferenceTarget(dataset, ref)
			if target == "" {
				continue
			}
			if _, ok := nodeSet[target]; !ok {
				continue
			}
			edges = append(edges, model.GraphEdge{
				Source:       r.ID,
				Target:       target,
				RelationType: ref.Relation,
				Label:        edgeLabelForReference(r, ref),
			})
		}

		for _, owner := range r.OwnerRefs {
			target := resolveOwner(dataset, r.Namespace, owner)
			if target == "" {
				continue
			}
			if _, ok := nodeSet[target]; !ok {
				continue
			}
			edges = append(edges, model.GraphEdge{Source: r.ID, Target: target, RelationType: "owns"})
		}
	}

	edges = append(edges, inferServiceEdges(dataset, nodeSet)...)
	edges = dedupeEdges(edges)
	edges = filterEdges(edges, filters)

	groups := []model.GraphGroup{}
	if filters.GroupBy == "namespace" || filters.GroupBy == "kind" {
		nodes, edges, groups = applyGrouping(nodes, edges, filters)
		kindStats = recomputeKindStats(nodes)
	}

	if filters.FocusID != "" {
		nodes, edges = applyFocus(nodes, edges, filters.FocusID)
		kindStats = recomputeKindStats(nodes)
	}

	sort.Slice(nodes, func(i, j int) bool { return nodes[i].ID < nodes[j].ID })
	sort.Slice(edges, func(i, j int) bool {
		if edges[i].Source != edges[j].Source {
			return edges[i].Source < edges[j].Source
		}
		if edges[i].Target != edges[j].Target {
			return edges[i].Target < edges[j].Target
		}
		if edges[i].RelationType != edges[j].RelationType {
			return edges[i].RelationType < edges[j].RelationType
		}
		return edges[i].Label < edges[j].Label
	})

	return model.GraphResponse{
		Nodes:  nodes,
		Edges:  edges,
		Groups: groups,
		Stats: model.GraphStats{
			TotalNodes: len(nodes),
			TotalEdges: len(edges),
			Kinds:      kindStats,
		},
	}
}

func applyGrouping(nodes []model.GraphNode, edges []model.GraphEdge, filters Filters) ([]model.GraphNode, []model.GraphEdge, []model.GraphGroup) {
	groupMembers := map[string][]model.GraphNode{}
	for _, n := range nodes {
		key := groupKey(n, filters.GroupBy)
		groupMembers[key] = append(groupMembers[key], n)
	}

	groups := make([]model.GraphGroup, 0, len(groupMembers))
	for key, members := range groupMembers {
		groups = append(groups, model.GraphGroup{
			Key:       key,
			Label:     key,
			Mode:      filters.GroupBy,
			Count:     len(members),
			Collapsed: filters.CollapsedGroups[key],
		})
	}
	sort.Slice(groups, func(i, j int) bool { return groups[i].Key < groups[j].Key })

	visibleNodes := make(map[string]model.GraphNode)
	nodeToVisible := map[string]string{}

	for key, members := range groupMembers {
		collapsed := filters.CollapsedGroups[key]
		if collapsed {
			groupID := groupNodeID(filters.GroupBy, key)
			health := "healthy"
			sensitive := false
			for _, m := range members {
				if m.HealthHint == "warning" {
					health = "warning"
				}
				if m.IsSensitive {
					sensitive = true
				}
				nodeToVisible[m.ID] = groupID
			}
			visibleNodes[groupID] = model.GraphNode{
				ID:          groupID,
				Kind:        strings.Title(filters.GroupBy) + "Group",
				Name:        key,
				Namespace:   "",
				HealthHint:  health,
				IsSensitive: sensitive,
				IsGroup:     true,
				GroupBy:     filters.GroupBy,
				GroupKey:    key,
				MemberCount: len(members),
			}
			continue
		}

		for _, m := range members {
			nodeToVisible[m.ID] = m.ID
			visibleNodes[m.ID] = m
		}
	}

	outNodes := make([]model.GraphNode, 0, len(visibleNodes))
	for _, n := range visibleNodes {
		outNodes = append(outNodes, n)
	}

	outEdges := make([]model.GraphEdge, 0, len(edges))
	for _, e := range edges {
		src := nodeToVisible[e.Source]
		tgt := nodeToVisible[e.Target]
		if src == "" || tgt == "" || src == tgt {
			continue
		}
		outEdges = append(outEdges, model.GraphEdge{
			Source:       src,
			Target:       tgt,
			RelationType: e.RelationType,
			Label:        e.Label,
		})
	}
	outEdges = dedupeEdges(outEdges)

	return outNodes, outEdges, groups
}

func groupKey(node model.GraphNode, mode string) string {
	switch mode {
	case "namespace":
		if node.Namespace == "" {
			return "(cluster-scoped)"
		}
		return node.Namespace
	case "kind":
		return node.Kind
	default:
		return ""
	}
}

func groupNodeID(mode, key string) string {
	return "group/" + mode + "/" + key
}

func resolveReferenceTarget(dataset *model.Dataset, ref model.ResourceReference) string {
	ns := ref.Namespace
	if ns == "" {
		ns = "default"
	}
	candidate := parser.ResourceID(ns, ref.Kind, ref.Name)
	if _, ok := dataset.Resources[candidate]; ok {
		return candidate
	}
	candidate = parser.ResourceID("", ref.Kind, ref.Name)
	if _, ok := dataset.Resources[candidate]; ok {
		return candidate
	}

	matches := make([]string, 0)
	for _, r := range dataset.Resources {
		if r.Kind == ref.Kind && r.Name == ref.Name {
			matches = append(matches, r.ID)
		}
	}
	if len(matches) == 1 {
		return matches[0]
	}
	return ""
}

func resolveOwner(dataset *model.Dataset, namespace string, owner model.OwnerReference) string {
	candidate := parser.ResourceID(namespace, owner.Kind, owner.Name)
	if _, ok := dataset.Resources[candidate]; ok {
		return candidate
	}
	candidate = parser.ResourceID("", owner.Kind, owner.Name)
	if _, ok := dataset.Resources[candidate]; ok {
		return candidate
	}
	return ""
}

func inferServiceEdges(dataset *model.Dataset, nodeSet map[string]struct{}) []model.GraphEdge {
	edges := []model.GraphEdge{}
	services := []*model.Resource{}
	workloads := []*model.Resource{}

	for _, r := range dataset.Resources {
		switch r.Kind {
		case "Service":
			services = append(services, r)
		case "Deployment", "StatefulSet", "DaemonSet":
			workloads = append(workloads, r)
		}
	}

	for _, svc := range services {
		if _, ok := nodeSet[svc.ID]; !ok {
			continue
		}
		if svc.WorkloadMeta == nil || len(svc.WorkloadMeta.ServiceSelectors) == 0 {
			continue
		}
		for _, wl := range workloads {
			if _, ok := nodeSet[wl.ID]; !ok {
				continue
			}
			if wl.WorkloadMeta == nil || len(wl.WorkloadMeta.PodTemplateLabels) == 0 {
				continue
			}
			if svc.Namespace != wl.Namespace {
				continue
			}
			if selectorsMatch(svc.WorkloadMeta.ServiceSelectors, wl.WorkloadMeta.PodTemplateLabels) {
				edges = append(edges, model.GraphEdge{
					Source:       svc.ID,
					Target:       wl.ID,
					RelationType: "selects",
					Label:        servicePortsLabel(svc),
				})
			}
		}
	}

	return edges
}

func selectorsMatch(selector, labels map[string]string) bool {
	for k, v := range selector {
		if labels[k] != v {
			return false
		}
	}
	return true
}

func matchesResourceFilters(node model.GraphNode, filters Filters) bool {
	if filters.Namespace != "" && node.Namespace != filters.Namespace {
		return false
	}
	if filters.Kind != "" && !strings.EqualFold(node.Kind, filters.Kind) {
		return false
	}
	if filters.Query != "" {
		needle := strings.ToLower(filters.Query)
		hay := strings.ToLower(node.ID + " " + node.Kind + " " + node.Name + " " + node.Namespace)
		if !strings.Contains(hay, needle) {
			return false
		}
	}
	return true
}

func dedupeEdges(edges []model.GraphEdge) []model.GraphEdge {
	seen := map[string]struct{}{}
	out := make([]model.GraphEdge, 0, len(edges))
	for _, edge := range edges {
		key := edge.Source + "|" + edge.Target + "|" + edge.RelationType + "|" + edge.Label
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		out = append(out, edge)
	}
	return out
}

func filterEdges(edges []model.GraphEdge, filters Filters) []model.GraphEdge {
	if len(filters.Relations) == 0 {
		return edges
	}
	out := make([]model.GraphEdge, 0, len(edges))
	for _, edge := range edges {
		if filters.Relations[edge.RelationType] {
			out = append(out, edge)
		}
	}
	return out
}

func applyFocus(nodes []model.GraphNode, edges []model.GraphEdge, focusID string) ([]model.GraphNode, []model.GraphEdge) {
	neighbor := map[string]struct{}{focusID: {}}
	for _, e := range edges {
		if e.Source == focusID {
			neighbor[e.Target] = struct{}{}
		}
		if e.Target == focusID {
			neighbor[e.Source] = struct{}{}
		}
	}

	fNodes := make([]model.GraphNode, 0, len(neighbor))
	for _, n := range nodes {
		if _, ok := neighbor[n.ID]; ok {
			fNodes = append(fNodes, n)
		}
	}

	fEdges := make([]model.GraphEdge, 0, len(edges))
	for _, e := range edges {
		_, src := neighbor[e.Source]
		_, dst := neighbor[e.Target]
		if src && dst {
			fEdges = append(fEdges, e)
		}
	}

	return fNodes, fEdges
}

func recomputeKindStats(nodes []model.GraphNode) map[string]int {
	out := map[string]int{}
	for _, n := range nodes {
		out[n.Kind]++
	}
	return out
}

func edgeLabelForReference(source *model.Resource, ref model.ResourceReference) string {
	if source == nil {
		return ""
	}

	switch source.Kind {
	case "Ingress":
		if ref.Relation != "routesTo" {
			return ""
		}
		return ingressBackendPortLabel(source, ref.Name)
	default:
		return ""
	}
}

func ingressBackendPortLabel(ingress *model.Resource, serviceName string) string {
	spec, _ := ingress.Raw["spec"].(map[string]any)
	if spec == nil {
		return ""
	}

	if backend, ok := spec["defaultBackend"].(map[string]any); ok {
		if name, label := ingressBackendServiceAndPort(backend); name == serviceName {
			return label
		}
	}

	rules, _ := spec["rules"].([]any)
	for _, r := range rules {
		rule, ok := r.(map[string]any)
		if !ok {
			continue
		}
		httpSpec, _ := rule["http"].(map[string]any)
		paths, _ := httpSpec["paths"].([]any)
		for _, p := range paths {
			pathEntry, ok := p.(map[string]any)
			if !ok {
				continue
			}
			backend, _ := pathEntry["backend"].(map[string]any)
			if name, label := ingressBackendServiceAndPort(backend); name == serviceName {
				return label
			}
		}
	}
	return ""
}

func ingressBackendServiceAndPort(backend map[string]any) (string, string) {
	service, _ := backend["service"].(map[string]any)
	if service == nil {
		return "", ""
	}

	name, _ := service["name"].(string)
	portMap, _ := service["port"].(map[string]any)
	if portMap == nil {
		return name, ""
	}
	if number, ok := portMap["number"]; ok {
		return name, fmt.Sprintf(":%v", number)
	}
	if pname, ok := portMap["name"].(string); ok && pname != "" {
		return name, ":" + pname
	}
	return name, ""
}

func servicePortsLabel(service *model.Resource) string {
	spec, _ := service.Raw["spec"].(map[string]any)
	if spec == nil {
		return ""
	}
	ports, _ := spec["ports"].([]any)
	if len(ports) == 0 {
		return ""
	}

	parts := make([]string, 0, len(ports))
	for _, p := range ports {
		port, ok := p.(map[string]any)
		if !ok {
			continue
		}
		protocol, _ := port["protocol"].(string)
		if protocol == "" {
			protocol = "TCP"
		}
		external, hasPort := port["port"]
		if !hasPort {
			continue
		}
		target, hasTarget := port["targetPort"]
		if hasTarget {
			parts = append(parts, fmt.Sprintf("%s %v->%v", protocol, external, target))
			continue
		}
		parts = append(parts, fmt.Sprintf("%s %v", protocol, external))
	}

	switch len(parts) {
	case 0:
		return ""
	case 1:
		return parts[0]
	default:
		return parts[0] + " +" + fmt.Sprintf("%d", len(parts)-1)
	}
}