mermaid-go/pkg/renderer/class.go

// Package renderer provides rendering functionality for class diagrams
package renderer

import (
	"fmt"
	"strings"

	"mermaid-go/pkg/ast"
)

// ClassRenderer renders class diagrams back to mermaid syntax
type ClassRenderer struct{}

// NewClassRenderer creates a new class renderer
func NewClassRenderer() *ClassRenderer {
	return &ClassRenderer{}
}

// Render renders a class diagram to mermaid syntax
func (r *ClassRenderer) Render(diagram *ast.ClassDiagram) (string, error) {
	var builder strings.Builder

	// Start with diagram declaration
	builder.WriteString("classDiagram\n")

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

	// Add direction if present
	if diagram.Direction != "" {
		builder.WriteString(fmt.Sprintf("    direction %s\n", diagram.Direction))
	}

	// Render classes
	for _, class := range diagram.Classes {
		builder.WriteString("    class ")
		builder.WriteString(class.ID)

		// Render class body if it has members or methods
		if len(class.Members) > 0 || len(class.Methods) > 0 || len(class.Annotations) > 0 {
			builder.WriteString(" {\n")

			// Render annotations first
			for _, annotation := range class.Annotations {
				builder.WriteString("        <<")
				builder.WriteString(annotation)
				builder.WriteString(">>\n")
			}

			// Render members
			for _, member := range class.Members {
				builder.WriteString("        ")
				builder.WriteString(string(member.Visibility))
				builder.WriteString(member.Name)
				if member.Type != "" {
					builder.WriteString(" ")
					// Handle generics in member type
					if len(member.Generics) > 0 {
						builder.WriteString(r.renderGenericType(member.Type, member.Generics))
					} else {
						builder.WriteString(member.Type)
					}
				}
				if member.Classifier != nil {
					builder.WriteString(" ")
					builder.WriteString(*member.Classifier)
				}
				builder.WriteString("\n")
			}

			// Render methods
			for _, method := range class.Methods {
				builder.WriteString("        ")
				builder.WriteString(string(method.Visibility))
				builder.WriteString(method.Name)
				builder.WriteString("(")

				for i, param := range method.Parameters {
					if i > 0 {
						builder.WriteString(", ")
					}
					builder.WriteString(param)
				}

				builder.WriteString(")")
				if method.Type != "" {
					builder.WriteString(" ")
					// Handle generics in return type
					if len(method.Generics) > 0 {
						builder.WriteString(r.renderGenericType(method.Type, method.Generics))
					} else {
						builder.WriteString(method.Type)
					}
				}
				if method.Classifier != nil {
					builder.WriteString(" ")
					builder.WriteString(*method.Classifier)
				}
				builder.WriteString("\n")
			}

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

		// Render links
		if class.Link != nil {
			builder.WriteString("    link ")
			builder.WriteString(class.ID)
			builder.WriteString(" \"")
			builder.WriteString(*class.Link)
			builder.WriteString("\"")
			if class.LinkTarget != nil {
				builder.WriteString(" ")
				builder.WriteString(*class.LinkTarget)
			}
			builder.WriteString("\n")
		}

		// Render tooltip
		if class.Tooltip != nil {
			builder.WriteString("    ")
			builder.WriteString(class.ID)
			builder.WriteString(" : ")
			builder.WriteString(*class.Tooltip)
			builder.WriteString("\n")
		}
	}

	// Render relations
	for _, relation := range diagram.Relations {
		builder.WriteString("    ")
		builder.WriteString(relation.From)
		builder.WriteString(" ")

		// Render relation type
		switch relation.Type {
		case ast.RelationInheritance:
			builder.WriteString("<|--")
		case ast.RelationComposition:
			builder.WriteString("*--")
		case ast.RelationAggregation:
			builder.WriteString("o--")
		case ast.RelationAssociation:
			builder.WriteString("-->")
		case ast.RelationRealization:
			builder.WriteString("..|>")
		case ast.RelationDependency:
			builder.WriteString("..>")
		default:
			builder.WriteString("-->") // Default to association
		}

		builder.WriteString(" ")
		builder.WriteString(relation.To)

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

		// Add cardinality if present
		if relation.Cardinality != nil {
			if relation.Cardinality.From != "" {
				builder.WriteString(" \"")
				builder.WriteString(relation.Cardinality.From)
				builder.WriteString("\"")
			}
			if relation.Cardinality.To != "" {
				builder.WriteString(" \"")
				builder.WriteString(relation.Cardinality.To)
				builder.WriteString("\"")
			}
		}

		builder.WriteString("\n")
	}

	// Render class definitions
	for _, classDef := range diagram.ClassDefs {
		builder.WriteString("    classDef ")
		builder.WriteString(classDef.ID)

		for _, style := range classDef.Styles {
			builder.WriteString(" ")
			builder.WriteString(style)
		}

		builder.WriteString("\n")
	}

	// Render notes
	for _, note := range diagram.Notes {
		if note.ForClass != nil {
			builder.WriteString(fmt.Sprintf("    note for %s \"%s\"\n", *note.ForClass, note.Text))
		} else {
			builder.WriteString(fmt.Sprintf("    note \"%s\"\n", note.Text))
		}
	}

	return builder.String(), nil
}

// renderGenericType renders a type with its generic parameters
func (r *ClassRenderer) renderGenericType(baseType string, generics []*ast.Generic) string {
	if len(generics) == 0 {
		return baseType
	}

	result := "~" + baseType
	if len(generics) > 0 {
		result += "<"
		for i, generic := range generics {
			if i > 0 {
				result += ", "
			}
			if len(generic.Arguments) > 0 {
				result += r.renderGenericType(generic.Name, generic.Arguments)
			} else {
				result += generic.Name
			}
		}
		result += ">"
	}
	result += "~"
	return result
}