relspecgo/pkg/writers/pgsql/migration_writer.go

package pgsql

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strings"

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

// MigrationWriter generates differential migration SQL scripts
type MigrationWriter struct {
	options *writers.WriterOptions
	writer  io.Writer
}

// MigrationScript represents a single migration script with priority and sequence
type MigrationScript struct {
	ObjectName string
	ObjectType string
	Schema     string
	Priority   int
	Sequence   int
	Body       string
}

// NewMigrationWriter creates a new migration writer
func NewMigrationWriter(options *writers.WriterOptions) *MigrationWriter {
	return &MigrationWriter{
		options: options,
	}
}

// WriteMigration generates migration scripts by comparing model (desired) vs current (actual) database
func (w *MigrationWriter) WriteMigration(model *models.Database, current *models.Database) error {
	var writer io.Writer
	var file *os.File
	var err error

	// Use existing writer if already set (for testing)
	if w.writer != nil {
		writer = w.writer
	} else if w.options.OutputPath != "" {
		file, err = os.Create(w.options.OutputPath)
		if err != nil {
			return fmt.Errorf("failed to create output file: %w", err)
		}
		defer file.Close()
		writer = file
	} else {
		writer = os.Stdout
	}

	w.writer = writer

	// Generate all migration scripts
	scripts := make([]MigrationScript, 0)

	// Process each schema in the model
	for _, modelSchema := range model.Schemas {
		// Find corresponding schema in current database
		var currentSchema *models.Schema
		for _, cs := range current.Schemas {
			if strings.EqualFold(cs.Name, modelSchema.Name) {
				currentSchema = cs
				break
			}
		}

		// Generate schema-level scripts
		schemaScripts := w.generateSchemaScripts(modelSchema, currentSchema)
		scripts = append(scripts, schemaScripts...)
	}

	// Sort scripts by priority and sequence
	sort.Slice(scripts, func(i, j int) bool {
		if scripts[i].Priority != scripts[j].Priority {
			return scripts[i].Priority < scripts[j].Priority
		}
		return scripts[i].Sequence < scripts[j].Sequence
	})

	// Write header
	fmt.Fprintf(w.writer, "-- PostgreSQL Migration Script\n")
	fmt.Fprintf(w.writer, "-- Generated by RelSpec\n")
	fmt.Fprintf(w.writer, "-- Source: %s -> %s\n\n", current.Name, model.Name)

	// Write scripts
	for _, script := range scripts {
		fmt.Fprintf(w.writer, "-- Priority: %d | Type: %s | Object: %s\n",
			script.Priority, script.ObjectType, script.ObjectName)
		fmt.Fprintf(w.writer, "%s\n\n", script.Body)
	}

	return nil
}

// generateSchemaScripts generates migration scripts for a schema
func (w *MigrationWriter) generateSchemaScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Phase 1: Drop constraints and indexes that changed (Priority 11-50)
	if current != nil {
		scripts = append(scripts, w.generateDropScripts(model, current)...)
	}

	// Phase 2: Rename tables and columns (Priority 60-90)
	if current != nil {
		scripts = append(scripts, w.generateRenameScripts(model, current)...)
	}

	// Phase 3: Create/Alter tables and columns (Priority 100-145)
	scripts = append(scripts, w.generateTableScripts(model, current)...)

	// Phase 4: Create indexes (Priority 160-180)
	scripts = append(scripts, w.generateIndexScripts(model, current)...)

	// Phase 5: Create foreign keys (Priority 195)
	scripts = append(scripts, w.generateForeignKeyScripts(model, current)...)

	// Phase 6: Add comments (Priority 200+)
	scripts = append(scripts, w.generateCommentScripts(model, current)...)

	return scripts
}

// generateDropScripts generates DROP scripts for removed/changed objects
func (w *MigrationWriter) generateDropScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Build map of model tables for quick lookup
	modelTables := make(map[string]*models.Table)
	for _, table := range model.Tables {
		modelTables[strings.ToLower(table.Name)] = table
	}

	// Find constraints to drop
	for _, currentTable := range current.Tables {
		modelTable, existsInModel := modelTables[strings.ToLower(currentTable.Name)]

		if !existsInModel {
			// Table will be dropped, skip individual constraint drops
			continue
		}

		// Check each constraint in current database
		for constraintName, currentConstraint := range currentTable.Constraints {
			// Check if constraint exists in model
			modelConstraint, existsInModel := modelTable.Constraints[constraintName]

			shouldDrop := false

			if !existsInModel {
				shouldDrop = true
			} else if !constraintsEqual(modelConstraint, currentConstraint) {
				// Constraint changed, drop and recreate
				shouldDrop = true
			}

			if shouldDrop {
				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", current.Name, currentTable.Name, constraintName),
					ObjectType: "drop constraint",
					Schema:     current.Name,
					Priority:   11,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"ALTER TABLE %s.%s DROP CONSTRAINT IF EXISTS %s;",
						current.Name, currentTable.Name, constraintName,
					),
				}
				scripts = append(scripts, script)
			}
		}

		// Check indexes
		for indexName, currentIndex := range currentTable.Indexes {
			modelIndex, existsInModel := modelTable.Indexes[indexName]

			shouldDrop := false

			if !existsInModel {
				shouldDrop = true
			} else if !indexesEqual(modelIndex, currentIndex) {
				shouldDrop = true
			}

			if shouldDrop {
				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", current.Name, currentTable.Name, indexName),
					ObjectType: "drop index",
					Schema:     current.Name,
					Priority:   20,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"DROP INDEX IF EXISTS %s.%s CASCADE;",
						current.Name, indexName,
					),
				}
				scripts = append(scripts, script)
			}
		}
	}

	return scripts
}

// generateRenameScripts generates RENAME scripts for renamed objects
func (w *MigrationWriter) generateRenameScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// For now, we don't attempt to detect renames automatically
	// This would require GUID matching or other heuristics
	// Users would need to handle renames manually or through metadata

	// Suppress unused parameter warnings
	_ = model
	_ = current

	return scripts
}

// generateTableScripts generates CREATE/ALTER TABLE scripts
func (w *MigrationWriter) generateTableScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Build map of current tables
	currentTables := make(map[string]*models.Table)
	if current != nil {
		for _, table := range current.Tables {
			currentTables[strings.ToLower(table.Name)] = table
		}
	}

	// Process each model table
	for _, modelTable := range model.Tables {
		currentTable, exists := currentTables[strings.ToLower(modelTable.Name)]

		if !exists {
			// Table doesn't exist, create it
			script := w.generateCreateTableScript(model, modelTable)
			scripts = append(scripts, script)
		} else {
			// Table exists, check for column changes
			alterScripts := w.generateAlterTableScripts(model, modelTable, currentTable)
			scripts = append(scripts, alterScripts...)
		}
	}

	return scripts
}

// generateCreateTableScript generates a CREATE TABLE script
func (w *MigrationWriter) generateCreateTableScript(schema *models.Schema, table *models.Table) MigrationScript {
	var body strings.Builder

	body.WriteString(fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s.%s (\n", schema.Name, table.Name))

	// Get sorted columns
	columns := getSortedColumns(table.Columns)
	columnDefs := make([]string, 0, len(columns))

	for _, col := range columns {
		colDef := fmt.Sprintf("  %s %s", col.Name, col.Type)

		// Add default value if present
		if col.Default != nil {
			colDef += fmt.Sprintf(" DEFAULT %v", col.Default)
		}

		// Add NOT NULL if needed
		if col.NotNull {
			colDef += " NOT NULL"
		}

		columnDefs = append(columnDefs, colDef)
	}

	body.WriteString(strings.Join(columnDefs, ",\n"))
	body.WriteString("\n);")

	return MigrationScript{
		ObjectName: fmt.Sprintf("%s.%s", schema.Name, table.Name),
		ObjectType: "create table",
		Schema:     schema.Name,
		Priority:   100,
		Sequence:   0,
		Body:       body.String(),
	}
}

// generateAlterTableScripts generates ALTER TABLE scripts for column changes
func (w *MigrationWriter) generateAlterTableScripts(schema *models.Schema, modelTable *models.Table, currentTable *models.Table) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Build map of current columns
	currentColumns := make(map[string]*models.Column)
	for name, col := range currentTable.Columns {
		currentColumns[strings.ToLower(name)] = col
	}

	// Check each model column
	for _, modelCol := range modelTable.Columns {
		currentCol, exists := currentColumns[strings.ToLower(modelCol.Name)]

		if !exists {
			// Column doesn't exist, add it
			script := MigrationScript{
				ObjectName: fmt.Sprintf("%s.%s.%s", schema.Name, modelTable.Name, modelCol.Name),
				ObjectType: "create column",
				Schema:     schema.Name,
				Priority:   120,
				Sequence:   len(scripts),
				Body: fmt.Sprintf(
					"ALTER TABLE %s.%s\n  ADD COLUMN IF NOT EXISTS %s %s%s%s;",
					schema.Name, modelTable.Name, modelCol.Name, modelCol.Type,
					func() string {
						if modelCol.Default != nil {
							return fmt.Sprintf(" DEFAULT %v", modelCol.Default)
						}
						return ""
					}(),
					func() string {
						if modelCol.NotNull {
							return " NOT NULL"
						}
						return ""
					}(),
				),
			}
			scripts = append(scripts, script)
		} else if !columnsEqual(modelCol, currentCol) {
			// Column exists but type or properties changed
			if modelCol.Type != currentCol.Type {
				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", schema.Name, modelTable.Name, modelCol.Name),
					ObjectType: "alter column type",
					Schema:     schema.Name,
					Priority:   120,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"ALTER TABLE %s.%s\n  ALTER COLUMN %s TYPE %s;",
						schema.Name, modelTable.Name, modelCol.Name, modelCol.Type,
					),
				}
				scripts = append(scripts, script)
			}

			// Check default value changes
			if fmt.Sprintf("%v", modelCol.Default) != fmt.Sprintf("%v", currentCol.Default) {
				if modelCol.Default != nil {
					script := MigrationScript{
						ObjectName: fmt.Sprintf("%s.%s.%s", schema.Name, modelTable.Name, modelCol.Name),
						ObjectType: "alter column default",
						Schema:     schema.Name,
						Priority:   145,
						Sequence:   len(scripts),
						Body: fmt.Sprintf(
							"ALTER TABLE %s.%s\n  ALTER COLUMN %s SET DEFAULT %v;",
							schema.Name, modelTable.Name, modelCol.Name, modelCol.Default,
						),
					}
					scripts = append(scripts, script)
				} else {
					script := MigrationScript{
						ObjectName: fmt.Sprintf("%s.%s.%s", schema.Name, modelTable.Name, modelCol.Name),
						ObjectType: "alter column default",
						Schema:     schema.Name,
						Priority:   145,
						Sequence:   len(scripts),
						Body: fmt.Sprintf(
							"ALTER TABLE %s.%s\n  ALTER COLUMN %s DROP DEFAULT;",
							schema.Name, modelTable.Name, modelCol.Name,
						),
					}
					scripts = append(scripts, script)
				}
			}
		}
	}

	return scripts
}

// generateIndexScripts generates CREATE INDEX scripts
func (w *MigrationWriter) generateIndexScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Build map of current tables
	currentTables := make(map[string]*models.Table)
	if current != nil {
		for _, table := range current.Tables {
			currentTables[strings.ToLower(table.Name)] = table
		}
	}

	// Process each model table
	for _, modelTable := range model.Tables {
		currentTable := currentTables[strings.ToLower(modelTable.Name)]

		// Process each index in model
		for indexName, modelIndex := range modelTable.Indexes {
			shouldCreate := true

			// Check if index exists in current
			if currentTable != nil {
				if currentIndex, exists := currentTable.Indexes[indexName]; exists {
					if indexesEqual(modelIndex, currentIndex) {
						shouldCreate = false
					}
				}
			}

			if shouldCreate {
				unique := ""
				if modelIndex.Unique {
					unique = "UNIQUE "
				}

				indexType := "btree"
				if modelIndex.Type != "" {
					indexType = modelIndex.Type
				}

				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", model.Name, modelTable.Name, indexName),
					ObjectType: "create index",
					Schema:     model.Name,
					Priority:   180,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"CREATE %sINDEX IF NOT EXISTS %s\n  ON %s.%s USING %s (%s);",
						unique, indexName, model.Name, modelTable.Name, indexType,
						strings.Join(modelIndex.Columns, ", "),
					),
				}
				scripts = append(scripts, script)
			}
		}

		// Add primary key constraint if it exists
		for constraintName, constraint := range modelTable.Constraints {
			if constraint.Type == models.PrimaryKeyConstraint {
				shouldCreate := true

				if currentTable != nil {
					if currentConstraint, exists := currentTable.Constraints[constraintName]; exists {
						if constraintsEqual(constraint, currentConstraint) {
							shouldCreate = false
						}
					}
				}

				if shouldCreate {
					script := MigrationScript{
						ObjectName: fmt.Sprintf("%s.%s.%s", model.Name, modelTable.Name, constraintName),
						ObjectType: "create primary key",
						Schema:     model.Name,
						Priority:   160,
						Sequence:   len(scripts),
						Body: fmt.Sprintf(
							"DO $$\nBEGIN\n  IF NOT EXISTS (\n"+
								"    SELECT 1 FROM information_schema.table_constraints\n"+
								"    WHERE table_schema = '%s'\n"+
								"      AND table_name = '%s'\n"+
								"      AND constraint_name = '%s'\n"+
								"  ) THEN\n"+
								"    ALTER TABLE %s.%s\n"+
								"      ADD CONSTRAINT %s PRIMARY KEY (%s);\n"+
								"  END IF;\n"+
								"END;\n$$;",
							model.Name, modelTable.Name, constraintName,
							model.Name, modelTable.Name, constraintName,
							strings.Join(constraint.Columns, ", "),
						),
					}
					scripts = append(scripts, script)
				}
			}
		}
	}

	return scripts
}

// generateForeignKeyScripts generates ADD CONSTRAINT FOREIGN KEY scripts
func (w *MigrationWriter) generateForeignKeyScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Build map of current tables
	currentTables := make(map[string]*models.Table)
	if current != nil {
		for _, table := range current.Tables {
			currentTables[strings.ToLower(table.Name)] = table
		}
	}

	// Process each model table
	for _, modelTable := range model.Tables {
		currentTable := currentTables[strings.ToLower(modelTable.Name)]

		// Process each constraint
		for constraintName, constraint := range modelTable.Constraints {
			if constraint.Type != models.ForeignKeyConstraint {
				continue
			}

			shouldCreate := true

			// Check if constraint exists in current
			if currentTable != nil {
				if currentConstraint, exists := currentTable.Constraints[constraintName]; exists {
					if constraintsEqual(constraint, currentConstraint) {
						shouldCreate = false
					}
				}
			}

			if shouldCreate {
				onDelete := "NO ACTION"
				if constraint.OnDelete != "" {
					onDelete = strings.ToUpper(constraint.OnDelete)
				}

				onUpdate := "NO ACTION"
				if constraint.OnUpdate != "" {
					onUpdate = strings.ToUpper(constraint.OnUpdate)
				}

				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", model.Name, modelTable.Name, constraintName),
					ObjectType: "create foreign key",
					Schema:     model.Name,
					Priority:   195,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"ALTER TABLE %s.%s\n"+
							"  DROP CONSTRAINT IF EXISTS %s;\n\n"+
							"ALTER TABLE %s.%s\n"+
							"  ADD CONSTRAINT %s\n"+
							"  FOREIGN KEY (%s)\n"+
							"  REFERENCES %s.%s (%s)\n"+
							"  ON DELETE %s\n"+
							"  ON UPDATE %s\n"+
							"  DEFERRABLE;",
						model.Name, modelTable.Name, constraintName,
						model.Name, modelTable.Name, constraintName,
						strings.Join(constraint.Columns, ", "),
						constraint.ReferencedSchema, constraint.ReferencedTable,
						strings.Join(constraint.ReferencedColumns, ", "),
						onDelete, onUpdate,
					),
				}
				scripts = append(scripts, script)
			}
		}
	}

	return scripts
}

// generateCommentScripts generates COMMENT ON scripts
func (w *MigrationWriter) generateCommentScripts(model *models.Schema, current *models.Schema) []MigrationScript {
	scripts := make([]MigrationScript, 0)

	// Suppress unused parameter warning (current not used yet, could be used for diffing)
	_ = current

	// Process each model table
	for _, modelTable := range model.Tables {
		// Table comment
		if modelTable.Description != "" {
			script := MigrationScript{
				ObjectName: fmt.Sprintf("%s.%s", model.Name, modelTable.Name),
				ObjectType: "comment on table",
				Schema:     model.Name,
				Priority:   200,
				Sequence:   len(scripts),
				Body: fmt.Sprintf(
					"COMMENT ON TABLE %s.%s IS '%s';",
					model.Name, modelTable.Name, escapeQuote(modelTable.Description),
				),
			}
			scripts = append(scripts, script)
		}

		// Column comments
		for _, col := range modelTable.Columns {
			if col.Description != "" {
				script := MigrationScript{
					ObjectName: fmt.Sprintf("%s.%s.%s", model.Name, modelTable.Name, col.Name),
					ObjectType: "comment on column",
					Schema:     model.Name,
					Priority:   200,
					Sequence:   len(scripts),
					Body: fmt.Sprintf(
						"COMMENT ON COLUMN %s.%s.%s IS '%s';",
						model.Name, modelTable.Name, col.Name, escapeQuote(col.Description),
					),
				}
				scripts = append(scripts, script)
			}
		}
	}

	return scripts
}

// Comparison helper functions

func constraintsEqual(a, b *models.Constraint) bool {
	if a.Type != b.Type {
		return false
	}
	if len(a.Columns) != len(b.Columns) {
		return false
	}
	for i := range a.Columns {
		if !strings.EqualFold(a.Columns[i], b.Columns[i]) {
			return false
		}
	}
	if a.Type == models.ForeignKeyConstraint {
		if a.ReferencedTable != b.ReferencedTable || a.ReferencedSchema != b.ReferencedSchema {
			return false
		}
		if len(a.ReferencedColumns) != len(b.ReferencedColumns) {
			return false
		}
		for i := range a.ReferencedColumns {
			if !strings.EqualFold(a.ReferencedColumns[i], b.ReferencedColumns[i]) {
				return false
			}
		}
	}
	return true
}

func indexesEqual(a, b *models.Index) bool {
	if a.Unique != b.Unique {
		return false
	}
	if len(a.Columns) != len(b.Columns) {
		return false
	}
	for i := range a.Columns {
		if !strings.EqualFold(a.Columns[i], b.Columns[i]) {
			return false
		}
	}
	return true
}

func columnsEqual(a, b *models.Column) bool {
	if a.Type != b.Type {
		return false
	}
	if a.NotNull != b.NotNull {
		return false
	}
	if fmt.Sprintf("%v", a.Default) != fmt.Sprintf("%v", b.Default) {
		return false
	}
	return true
}