mermaid-go/pkg/renderer/er.go

// Package renderer provides ER diagram rendering
package renderer

import (
	"fmt"
	"strings"

	"mermaid-go/pkg/ast"
)

// ERRenderer renders ER diagrams back to mermaid syntax
type ERRenderer struct{}

// NewERRenderer creates a new ER renderer
func NewERRenderer() *ERRenderer {
	return &ERRenderer{}
}

// Render renders an ER diagram to mermaid syntax
func (r *ERRenderer) Render(diagram *ast.ERDiagram) (string, error) {
	var builder strings.Builder

	builder.WriteString("erDiagram\n")

	// Render title if present
	if diagram.Title != nil {
		builder.WriteString(fmt.Sprintf("    title %s\n", *diagram.Title))
	}

	// Render entities with attributes
	for _, entity := range diagram.Entities {
		if len(entity.Attributes) > 0 {
			builder.WriteString(fmt.Sprintf("    %s {\n", r.quoteIfNeeded(entity.Name)))
			for _, attr := range entity.Attributes {
				builder.WriteString("        ")
				builder.WriteString(attr.Type)
				builder.WriteString(" ")
				builder.WriteString(attr.Name)

				// Add key constraint if present
				if attr.Key != nil {
					switch *attr.Key {
					case ast.ERKeyPrimary:
						builder.WriteString(" PK")
					case ast.ERKeyForeign:
						builder.WriteString(" FK")
					case ast.ERKeyUnique:
						builder.WriteString(" UK")
					}
				}

				// Add comment if present
				if attr.Comment != nil {
					builder.WriteString(fmt.Sprintf(" \"%s\"", *attr.Comment))
				}

				builder.WriteString("\n")
			}
			builder.WriteString("    }\n")
		}
	}

	// Render relationships
	for _, relation := range diagram.Relations {
		builder.WriteString("    ")
		builder.WriteString(r.quoteIfNeeded(relation.From))
		builder.WriteString(" ")
		builder.WriteString(r.renderRelationType(relation.Type))
		builder.WriteString(" ")
		builder.WriteString(r.quoteIfNeeded(relation.To))

		// Add label if present
		if relation.Label != nil {
			builder.WriteString(" : ")
			builder.WriteString(*relation.Label)
		}

		builder.WriteString("\n")
	}

	return builder.String(), nil
}

// renderRelationType converts relation type to mermaid syntax
func (r *ERRenderer) renderRelationType(relType ast.ERRelationType) string {
	switch relType {
	case ast.ERRelationOneToOne:
		return "||--||"
	case ast.ERRelationOneToMany:
		return "||--o{"
	case ast.ERRelationManyToOne:
		return "}o--||"
	case ast.ERRelationManyToMany:
		return "}o--o{"
	default:
		return "||--o{" // Default to one-to-many
	}
}

// quoteIfNeeded adds quotes around entity names if they contain spaces
func (r *ERRenderer) quoteIfNeeded(name string) string {
	if strings.Contains(name, " ") {
		return fmt.Sprintf("\"%s\"", name)
	}
	return name
}