mermaid-go/pkg/detector/detector.go

package detector

import (
	"regexp"
	"strings"
)

// DiagramType represents the type of diagram
type DiagramType string

const (
	DiagramTypeUnknown      DiagramType = "unknown"
	DiagramTypeSequence     DiagramType = "sequence"
	DiagramTypeState        DiagramType = "state"
	DiagramTypeStateV2      DiagramType = "state-v2"
	DiagramTypeFlowchart    DiagramType = "flowchart"
	DiagramTypeClass        DiagramType = "class"
	DiagramTypeER           DiagramType = "er"
	DiagramTypeGantt        DiagramType = "gantt"
	DiagramTypePie          DiagramType = "pie"
	DiagramTypeGit          DiagramType = "git"
	DiagramTypeUserJourney  DiagramType = "user-journey"
	DiagramTypeTimeline     DiagramType = "timeline"
	DiagramTypeMindmap      DiagramType = "mindmap"
	DiagramTypeKanban       DiagramType = "kanban"
	DiagramTypeSankey       DiagramType = "sankey"
	DiagramTypeXYChart      DiagramType = "xychart"
	DiagramTypeQuadrant     DiagramType = "quadrant-chart"
	DiagramTypeRequirement  DiagramType = "requirement"
	DiagramTypeBlock        DiagramType = "block"
	DiagramTypeC4           DiagramType = "c4"
	DiagramTypeArchitecture DiagramType = "architecture"
	DiagramTypeRadar        DiagramType = "radar"
	DiagramTypeTreemap      DiagramType = "treemap"
	DiagramTypePacket       DiagramType = "packet"
	DiagramTypeInfo         DiagramType = "info"
)

// DetectorFunc is a function that detects if text matches a specific diagram type
type DetectorFunc func(text string) bool

// DiagramDetector represents a diagram detector with its type and detection function
type DiagramDetector struct {
	Type     DiagramType
	Detector DetectorFunc
	Priority int // Higher priority = checked first
}

// DetectorRegistry manages all diagram detectors
type DetectorRegistry struct {
	detectors []DiagramDetector
}

// NewDetectorRegistry creates a new detector registry with default detectors
func NewDetectorRegistry() *DetectorRegistry {
	registry := &DetectorRegistry{
		detectors: make([]DiagramDetector, 0),
	}

	// Register default detectors in order of specificity (most specific first)
	registry.Register(sequenceDetector())
	registry.Register(stateV2Detector())
	registry.Register(stateDetector())
	registry.Register(flowchartDetector())
	registry.Register(classDetector())
	registry.Register(erDetector())
	registry.Register(ganttDetector())
	registry.Register(pieDetector())
	registry.Register(gitDetector())
	registry.Register(userJourneyDetector())
	registry.Register(timelineDetector())
	registry.Register(mindmapDetector())
	registry.Register(kanbanDetector())
	registry.Register(sankeyDetector())
	registry.Register(xyChartDetector())
	registry.Register(quadrantDetector())
	registry.Register(requirementDetector())
	registry.Register(blockDetector())
	registry.Register(c4Detector())
	registry.Register(architectureDetector())
	registry.Register(radarDetector())
	registry.Register(treemapDetector())
	registry.Register(packetDetector())
	registry.Register(infoDetector())

	return registry
}

// Register adds a diagram detector to the registry
func (r *DetectorRegistry) Register(detector DiagramDetector) {
	r.detectors = append(r.detectors, detector)
	// Sort by priority (higher first)
	for i := len(r.detectors) - 1; i > 0; i-- {
		if r.detectors[i].Priority > r.detectors[i-1].Priority {
			r.detectors[i], r.detectors[i-1] = r.detectors[i-1], r.detectors[i]
		} else {
			break
		}
	}
}

// DetectDiagramType detects the type of diagram from the given text
func (r *DetectorRegistry) DetectDiagramType(text string) DiagramType {
	// Clean the text by removing comments and directives
	cleanText := cleanText(text)

	for _, detector := range r.detectors {
		if detector.Detector(cleanText) {
			return detector.Type
		}
	}

	return DiagramTypeUnknown
}

// cleanText removes comments, directives, and front matter from text
func cleanText(text string) string {
	// Remove front matter (YAML between --- markers)
	frontMatterRegex := regexp.MustCompile(`(?s)^---\s*\n(.*?)\n---\s*\n+`)
	text = frontMatterRegex.ReplaceAllString(text, "")

	// Remove directives (%%{...}%%)
	directiveRegex := regexp.MustCompile(`%%\{[^}]*\}%%`)
	text = directiveRegex.ReplaceAllString(text, "")

	// Remove comments (%% ...) - this will leave empty lines
	commentRegex := regexp.MustCompile(`(?m)^\s*%%.*$`)
	text = commentRegex.ReplaceAllString(text, "")

	// Remove multiple consecutive newlines and normalize whitespace
	text = regexp.MustCompile(`\n\s*\n+`).ReplaceAllString(text, "\n")

	// Remove leading/trailing whitespace from each line
	lines := strings.Split(text, "\n")
	for i, line := range lines {
		lines[i] = strings.TrimRight(line, " \t")
	}
	text = strings.Join(lines, "\n")

	// Remove leading/trailing whitespace from entire text
	text = strings.TrimSpace(text)

	return text
}

// Individual detector functions

func sequenceDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeSequence,
		Priority: 100,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*sequenceDiagram`, text)
			return matched
		},
	}
}

func stateV2Detector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeStateV2,
		Priority: 95,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*stateDiagram-v2`, text)
			return matched
		},
	}
}

func stateDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeState,
		Priority: 90,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*stateDiagram`, text)
			return matched
		},
	}
}

func flowchartDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeFlowchart,
		Priority: 85,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*(?:flowchart|graph)`, text)
			return matched
		},
	}
}

func classDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeClass,
		Priority: 80,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*classDiagram`, text)
			return matched
		},
	}
}

func erDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeER,
		Priority: 75,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*erDiagram`, text)
			return matched
		},
	}
}

func ganttDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeGantt,
		Priority: 70,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*gantt`, text)
			return matched
		},
	}
}

func pieDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypePie,
		Priority: 65,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*pie`, text)
			return matched
		},
	}
}

func gitDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeGit,
		Priority: 60,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*gitGraph`, text)
			return matched
		},
	}
}

func userJourneyDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeUserJourney,
		Priority: 55,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*journey`, text)
			return matched
		},
	}
}

func timelineDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeTimeline,
		Priority: 50,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*timeline`, text)
			return matched
		},
	}
}

func mindmapDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeMindmap,
		Priority: 45,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*mindmap`, text)
			return matched
		},
	}
}

func kanbanDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeKanban,
		Priority: 40,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*kanban`, text)
			return matched
		},
	}
}

func sankeyDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeSankey,
		Priority: 35,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*sankey-beta`, text)
			return matched
		},
	}
}

func xyChartDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeXYChart,
		Priority: 30,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*xychart-beta`, text)
			return matched
		},
	}
}

func quadrantDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeQuadrant,
		Priority: 25,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*quadrantChart`, text)
			return matched
		},
	}
}

func requirementDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeRequirement,
		Priority: 20,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*requirement`, text)
			return matched
		},
	}
}

func blockDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeBlock,
		Priority: 15,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*block-beta`, text)
			return matched
		},
	}
}

func c4Detector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeC4,
		Priority: 10,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*c4Context`, text)
			return matched
		},
	}
}

func architectureDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeArchitecture,
		Priority: 5,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*architecture`, text)
			return matched
		},
	}
}

func radarDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeRadar,
		Priority: 3,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*radar`, text)
			return matched
		},
	}
}

func treemapDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeTreemap,
		Priority: 2,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*treemap`, text)
			return matched
		},
	}
}

func packetDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypePacket,
		Priority: 1,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*packet`, text)
			return matched
		},
	}
}

func infoDetector() DiagramDetector {
	return DiagramDetector{
		Type:     DiagramTypeInfo,
		Priority: 0,
		Detector: func(text string) bool {
			matched, _ := regexp.MatchString(`(?i)^\s*info`, text)
			return matched
		},
	}
}