GSP/toolkit/internal/gsp/graph.go

package gsp

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"
)

func (p *Project) Graph(id string, depth int) Graph {
	if id != "" {
		flattened := p.Flatten(id, depth, Filter{})
		return p.graphForUnits(flattened.Units)
	}
	return p.graphForUnits(p.Units)
}

func (p *Project) graphForUnits(units []*Unit) Graph {
	nodeMap := map[string]GraphNode{}
	edgeMap := map[string]GraphEdge{}
	include := map[string]bool{}
	for _, unit := range units {
		include[unit.ID] = true
		nodeMap[unit.ID] = GraphNode{ID: unit.ID, Title: unit.Title, Type: unit.Type, Resolution: unit.Resolution, File: unit.File}
	}
	for _, unit := range units {
		if unit.Refines != "" {
			addEdge(edgeMap, unit.ID, unit.Refines, "refines")
			if !include[unit.Refines] {
				nodeMap[unit.Refines] = GraphNode{ID: unit.Refines, Missing: p.ByID[unit.Refines] == nil}
			}
		}
		for _, rel := range unit.With {
			addEdge(edgeMap, unit.ID, rel.ID, "with")
			if !include[rel.ID] {
				nodeMap[rel.ID] = GraphNode{ID: rel.ID, Missing: p.ByID[rel.ID] == nil}
			}
		}
	}
	nodes := make([]GraphNode, 0, len(nodeMap))
	for _, node := range nodeMap {
		nodes = append(nodes, node)
	}
	sort.Slice(nodes, func(i, j int) bool {
		return nodes[i].ID < nodes[j].ID
	})
	edges := make([]GraphEdge, 0, len(edgeMap))
	for _, edge := range edgeMap {
		edges = append(edges, edge)
	}
	sort.Slice(edges, func(i, j int) bool {
		if edges[i].From == edges[j].From {
			if edges[i].To == edges[j].To {
				return edges[i].Kind < edges[j].Kind
			}
			return edges[i].To < edges[j].To
		}
		return edges[i].From < edges[j].From
	})
	return Graph{Nodes: nodes, Edges: edges}
}

func addEdge(edges map[string]GraphEdge, from, to, kind string) {
	key := from + "\x00" + to + "\x00" + kind
	edges[key] = GraphEdge{From: from, To: to, Kind: kind}
}

var mermaidIDPattern = regexp.MustCompile(`[^A-Za-z0-9_]`)

func (g Graph) Mermaid() string {
	var builder strings.Builder
	builder.WriteString("graph TD\n")
	if len(g.Nodes) == 0 {
		return builder.String()
	}
	for _, node := range g.Nodes {
		label := node.DisplayLabel()
		if node.Missing {
			label += " (missing)"
		}
		builder.WriteString(fmt.Sprintf("  %s[\"%s\"]\n", mermaidID(node.ID), escapeMermaid(label)))
	}
	for _, edge := range g.Edges {
		builder.WriteString(fmt.Sprintf("  %s -- %s --> %s\n", mermaidID(edge.From), edge.Kind, mermaidID(edge.To)))
	}
	return builder.String()
}

func (g Graph) Markdown() string {
	var builder strings.Builder
	builder.WriteString("# GSP Graph\n\n")
	builder.WriteString("```mermaid\n")
	builder.WriteString(g.Mermaid())
	builder.WriteString("```\n")
	return builder.String()
}

func (p *Project) ImpactGraph(impact ImpactResult) Graph {
	nodeMap := map[string]GraphNode{}
	if unit, ok := p.ByID[impact.Entry]; ok {
		nodeMap[impact.Entry] = GraphNode{ID: unit.ID, Title: unit.Title, Type: unit.Type, Resolution: unit.Resolution, File: unit.File}
	} else {
		nodeMap[impact.Entry] = GraphNode{ID: impact.Entry, Missing: true}
	}
	for _, entry := range append(append([]ImpactEntry{}, impact.Direct...), impact.Indirect...) {
		nodeMap[entry.ID] = GraphNode{ID: entry.ID, Title: entry.Title, Type: entry.Type, Resolution: entry.Resolution, File: entry.File}
		if _, ok := nodeMap[entry.Via]; !ok && entry.Via != "" {
			if unit, exists := p.ByID[entry.Via]; exists {
				nodeMap[entry.Via] = GraphNode{ID: unit.ID, Title: unit.Title, Type: unit.Type, Resolution: unit.Resolution, File: unit.File}
			} else {
				nodeMap[entry.Via] = GraphNode{ID: entry.Via, Missing: true}
			}
		}
	}
	nodes := make([]GraphNode, 0, len(nodeMap))
	for _, node := range nodeMap {
		nodes = append(nodes, node)
	}
	sort.Slice(nodes, func(i, j int) bool {
		return nodes[i].ID < nodes[j].ID
	})
	return Graph{Nodes: nodes, Edges: impact.Edges}
}

func (g Graph) Canvas() ([]byte, error) {
	canvas := canvasDocument{
		Nodes: make([]canvasNode, 0, len(g.Nodes)),
		Edges: make([]canvasEdge, 0, len(g.Edges)),
	}
	positions := g.canvasPositions()
	for _, node := range g.Nodes {
		pos := positions[node.ID]
		canvas.Nodes = append(canvas.Nodes, canvasNode{
			ID:     canvasNodeID(node.ID),
			Type:   "text",
			Text:   node.CanvasText(),
			X:      pos.X,
			Y:      pos.Y,
			Width:  280,
			Height: 140,
			Color:  node.CanvasColor(),
		})
	}
	for _, edge := range g.Edges {
		canvas.Edges = append(canvas.Edges, canvasEdge{
			ID:       canvasEdgeID(edge),
			FromNode: canvasNodeID(edge.From),
			FromSide: "right",
			ToNode:   canvasNodeID(edge.To),
			ToSide:   "left",
			Label:    edge.Kind,
		})
	}
	data, err := json.MarshalIndent(canvas, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(data, '\n'), nil
}

func (n GraphNode) DisplayLabel() string {
	title := strings.TrimSpace(n.Title)
	if title != "" {
		return title
	}
	return n.ID
}

func (n GraphNode) CanvasText() string {
	var lines []string
	lines = append(lines, "## "+n.DisplayLabel())
	if n.Title != "" && n.ID != "" {
		lines = append(lines, "`"+n.ID+"`")
	}
	if n.Type != "" {
		lines = append(lines, "type: "+n.Type)
	}
	if n.Resolution != "" {
		lines = append(lines, "resolution: "+n.Resolution)
	}
	if n.Missing {
		lines = append(lines, "missing: true")
	}
	return strings.Join(lines, "\n")
}

func (n GraphNode) CanvasColor() string {
	if n.Missing {
		return "1"
	}
	switch n.Resolution {
	case "L0", "L1":
		return "2"
	case "L2":
		return "3"
	case "L3":
		return "4"
	case "L4", "L5":
		return "5"
	default:
		return ""
	}
}

type canvasDocument struct {
	Nodes []canvasNode `json:"nodes"`
	Edges []canvasEdge `json:"edges"`
}

type canvasNode struct {
	ID     string `json:"id"`
	Type   string `json:"type"`
	Text   string `json:"text"`
	X      int    `json:"x"`
	Y      int    `json:"y"`
	Width  int    `json:"width"`
	Height int    `json:"height"`
	Color  string `json:"color,omitempty"`
}

type canvasEdge struct {
	ID       string `json:"id"`
	FromNode string `json:"fromNode"`
	FromSide string `json:"fromSide,omitempty"`
	ToNode   string `json:"toNode"`
	ToSide   string `json:"toSide,omitempty"`
	Label    string `json:"label,omitempty"`
}

type canvasPosition struct {
	X int
	Y int
}

func (g Graph) canvasPositions() map[string]canvasPosition {
	incoming := map[string]int{}
	outgoing := map[string][]string{}
	for _, node := range g.Nodes {
		incoming[node.ID] = 0
	}
	for _, edge := range g.Edges {
		if _, ok := incoming[edge.From]; !ok {
			incoming[edge.From] = 0
		}
		if _, ok := incoming[edge.To]; !ok {
			incoming[edge.To] = 0
		}
		incoming[edge.To]++
		outgoing[edge.From] = append(outgoing[edge.From], edge.To)
	}
	for id := range outgoing {
		sort.Strings(outgoing[id])
	}
	var queue []string
	for _, node := range g.Nodes {
		if incoming[node.ID] == 0 {
			queue = append(queue, node.ID)
		}
	}
	if len(queue) == 0 {
		for _, node := range g.Nodes {
			queue = append(queue, node.ID)
		}
	}
	sort.Strings(queue)
	depth := map[string]int{}
	for len(queue) > 0 {
		id := queue[0]
		queue = queue[1:]
		for _, to := range outgoing[id] {
			if depth[to] < depth[id]+1 {
				depth[to] = depth[id] + 1
			}
			incoming[to]--
			if incoming[to] == 0 {
				queue = append(queue, to)
				sort.Strings(queue)
			}
		}
	}
	columns := map[int][]string{}
	for _, node := range g.Nodes {
		columns[depth[node.ID]] = append(columns[depth[node.ID]], node.ID)
	}
	positions := map[string]canvasPosition{}
	var keys []int
	for key := range columns {
		keys = append(keys, key)
	}
	sort.Ints(keys)
	for _, key := range keys {
		sort.Strings(columns[key])
		for row, id := range columns[key] {
			positions[id] = canvasPosition{X: key * 360, Y: row * 190}
		}
	}
	return positions
}

func mermaidID(id string) string {
	value := mermaidIDPattern.ReplaceAllString(id, "_")
	if value == "" {
		return "node"
	}
	if value[0] >= '0' && value[0] <= '9' {
		value = "n_" + value
	}
	return value
}

func escapeMermaid(value string) string {
	value = strings.ReplaceAll(value, `"`, `\"`)
	return value
}

func canvasNodeID(id string) string {
	return "gsp_" + mermaidID(id)
}

func canvasEdgeID(edge GraphEdge) string {
	return "edge_" + mermaidID(edge.From) + "_" + mermaidID(edge.Kind) + "_" + mermaidID(edge.To)
}