relspecgo/pkg/writers/dbml/writer.go

package dbml

import (
	"fmt"
	"os"
	"strings"

	"git.warky.dev/wdevs/relspecgo/pkg/models"
	"git.warky.dev/wdevs/relspecgo/pkg/writers"
)

// Writer implements the writers.Writer interface for DBML format
type Writer struct {
	options *writers.WriterOptions
}

// NewWriter creates a new DBML writer with the given options
func NewWriter(options *writers.WriterOptions) *Writer {
	return &Writer{
		options: options,
	}
}

// WriteDatabase writes a Database model to DBML format
func (w *Writer) WriteDatabase(db *models.Database) error {
	content := w.databaseToDBML(db)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	// If no output path, print to stdout
	fmt.Print(content)
	return nil
}

// WriteSchema writes a Schema model to DBML format
func (w *Writer) WriteSchema(schema *models.Schema) error {
	content := w.schemaToDBML(schema)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	fmt.Print(content)
	return nil
}

// WriteTable writes a Table model to DBML format
func (w *Writer) WriteTable(table *models.Table) error {
	content := w.tableToDBML(table, table.Schema)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	fmt.Print(content)
	return nil
}

// databaseToDBML converts a Database to DBML format string
func (w *Writer) databaseToDBML(d *models.Database) string {
	var result string

	// Add database comment if exists
	if d.Description != "" {
		result += fmt.Sprintf("// %s\n", d.Description)
	}
	if d.Comment != "" {
		result += fmt.Sprintf("// %s\n", d.Comment)
	}
	if d.Description != "" || d.Comment != "" {
		result += "\n"
	}

	// Process each schema
	for _, schema := range d.Schemas {
		result += w.schemaToDBML(schema)
	}

	// Add relationships
	result += "\n// Relationships\n"
	for _, schema := range d.Schemas {
		for _, table := range schema.Tables {
			for _, constraint := range table.Constraints {
				if constraint.Type == models.ForeignKeyConstraint {
					result += w.constraintToDBML(constraint, schema.Name, table.Name)
				}
			}
		}
	}

	return result
}

// schemaToDBML converts a Schema to DBML format string
func (w *Writer) schemaToDBML(schema *models.Schema) string {
	var result string

	if schema.Description != "" {
		result += fmt.Sprintf("// Schema: %s - %s\n", schema.Name, schema.Description)
	}

	// Process tables
	for _, table := range schema.Tables {
		result += w.tableToDBML(table, schema.Name)
		result += "\n"
	}

	return result
}

// tableToDBML converts a Table to DBML format string
func (w *Writer) tableToDBML(t *models.Table, schemaName string) string {
	var result string

	// Table definition
	tableName := fmt.Sprintf("%s.%s", schemaName, t.Name)
	result += fmt.Sprintf("Table %s {\n", tableName)

	// Add columns
	for _, column := range t.Columns {
		result += fmt.Sprintf("  %s %s", column.Name, column.Type)

		// Add column attributes
		attrs := make([]string, 0)
		if column.IsPrimaryKey {
			attrs = append(attrs, "pk")
		}
		if column.NotNull && !column.IsPrimaryKey {
			attrs = append(attrs, "not null")
		}
		if column.AutoIncrement {
			attrs = append(attrs, "increment")
		}
		if column.Default != nil {
			attrs = append(attrs, fmt.Sprintf("default: %v", column.Default))
		}

		if len(attrs) > 0 {
			result += fmt.Sprintf(" [%s]", strings.Join(attrs, ", "))
		}

		if column.Comment != "" {
			result += fmt.Sprintf(" // %s", column.Comment)
		}
		result += "\n"
	}

	// Add indexes
	indexCount := 0
	for _, index := range t.Indexes {
		if indexCount == 0 {
			result += "\n  indexes {\n"
		}
		indexAttrs := make([]string, 0)
		if index.Unique {
			indexAttrs = append(indexAttrs, "unique")
		}
		if index.Name != "" {
			indexAttrs = append(indexAttrs, fmt.Sprintf("name: '%s'", index.Name))
		}
		if index.Type != "" {
			indexAttrs = append(indexAttrs, fmt.Sprintf("type: %s", index.Type))
		}

		result += fmt.Sprintf("    (%s)", strings.Join(index.Columns, ", "))
		if len(indexAttrs) > 0 {
			result += fmt.Sprintf(" [%s]", strings.Join(indexAttrs, ", "))
		}
		result += "\n"
		indexCount++
	}
	if indexCount > 0 {
		result += "  }\n"
	}

	// Add table note
	if t.Description != "" || t.Comment != "" {
		note := t.Description
		if note != "" && t.Comment != "" {
			note += " - "
		}
		note += t.Comment
		result += fmt.Sprintf("\n  Note: '%s'\n", note)
	}

	result += "}\n"
	return result
}

// constraintToDBML converts a Constraint to DBML format string
func (w *Writer) constraintToDBML(c *models.Constraint, schemaName, tableName string) string {
	if c.Type != models.ForeignKeyConstraint || c.ReferencedTable == "" {
		return ""
	}

	fromTable := fmt.Sprintf("%s.%s", schemaName, tableName)
	toTable := fmt.Sprintf("%s.%s", c.ReferencedSchema, c.ReferencedTable)

	// Determine relationship cardinality
	// For foreign keys, it's typically many-to-one
	relationship := ">"

	// Build from and to column references
	// For single columns: table.column
	// For multiple columns: table.(col1, col2)
	var fromRef, toRef string
	if len(c.Columns) == 1 {
		fromRef = fmt.Sprintf("%s.%s", fromTable, c.Columns[0])
	} else {
		fromRef = fmt.Sprintf("%s.(%s)", fromTable, strings.Join(c.Columns, ", "))
	}

	if len(c.ReferencedColumns) == 1 {
		toRef = fmt.Sprintf("%s.%s", toTable, c.ReferencedColumns[0])
	} else {
		toRef = fmt.Sprintf("%s.(%s)", toTable, strings.Join(c.ReferencedColumns, ", "))
	}

	result := fmt.Sprintf("Ref: %s %s %s", fromRef, relationship, toRef)

	// Add actions
	actions := make([]string, 0)
	if c.OnDelete != "" {
		actions = append(actions, fmt.Sprintf("ondelete: %s", c.OnDelete))
	}
	if c.OnUpdate != "" {
		actions = append(actions, fmt.Sprintf("onupdate: %s", c.OnUpdate))
	}
	if len(actions) > 0 {
		result += fmt.Sprintf(" [%s]", strings.Join(actions, ", "))
	}

	result += "\n"
	return result
}