relspecgo/pkg/writers/drizzle/template_data.go

package drizzle

import (
	"sort"

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

// TemplateData represents the data passed to the template for code generation
type TemplateData struct {
	Imports []string
	Enums   []*EnumData
	Tables  []*TableData
}

// EnumData represents an enum in the schema
type EnumData struct {
	Name       string   // Enum name (PascalCase)
	VarName    string   // Variable name for the enum (camelCase)
	Values     []string // Enum values
	ValuesStr  string   // Comma-separated quoted values for pgEnum()
	TypeUnion  string   // TypeScript union type (e.g., "'admin' | 'user' | 'guest'")
	SchemaName string   // Schema name
}

// TableData represents a table in the template
type TableData struct {
	Name              string        // Table variable name (camelCase, e.g., users)
	TableName         string        // Actual database table name (e.g., users)
	TypeName          string        // TypeScript type name (PascalCase, e.g., Users)
	Columns           []*ColumnData // Column definitions
	Indexes           []*IndexData  // Index definitions
	Comment           string        // Table comment
	SchemaName        string        // Schema name
	NeedsSQLTag       bool          // Whether we need to import 'sql' from drizzle-orm
	IndexColumnFields []string      // Column field names used in indexes (for destructuring)
}

// ColumnData represents a column in a table
type ColumnData struct {
	Name           string // Column name in database
	FieldName      string // Field name in TypeScript (camelCase)
	DrizzleChain   string // Complete Drizzle column chain (e.g., "integer('id').primaryKey()")
	TypeScriptType string // TypeScript type for interface (e.g., "string", "number | null")
	IsForeignKey   bool   // Whether this is a foreign key
	ReferencesLine string // The .references() line if FK
	Comment        string // Column comment
}

// IndexData represents an index definition
type IndexData struct {
	Name       string   // Index name
	Columns    []string // Column names
	IsUnique   bool     // Whether it's a unique index
	Definition string   // Complete index definition line
}

// NewTemplateData creates a new TemplateData
func NewTemplateData() *TemplateData {
	return &TemplateData{
		Imports: make([]string, 0),
		Enums:   make([]*EnumData, 0),
		Tables:  make([]*TableData, 0),
	}
}

// AddImport adds an import to the template data (deduplicates automatically)
func (td *TemplateData) AddImport(importLine string) {
	// Check if already exists
	for _, imp := range td.Imports {
		if imp == importLine {
			return
		}
	}
	td.Imports = append(td.Imports, importLine)
}

// AddEnum adds an enum to the template data
func (td *TemplateData) AddEnum(enum *EnumData) {
	td.Enums = append(td.Enums, enum)
}

// AddTable adds a table to the template data
func (td *TemplateData) AddTable(table *TableData) {
	td.Tables = append(td.Tables, table)
}

// FinalizeImports sorts imports
func (td *TemplateData) FinalizeImports() {
	sort.Strings(td.Imports)
}

// NewEnumData creates EnumData from a models.Enum
func NewEnumData(enum *models.Enum, tm *TypeMapper) *EnumData {
	// Keep enum name as-is (it should already be PascalCase from the source)
	enumName := enum.Name
	// Variable name is camelCase version
	varName := tm.ToCamelCase(enum.Name)

	// Format values as comma-separated quoted strings for pgEnum()
	quotedValues := make([]string, len(enum.Values))
	for i, v := range enum.Values {
		quotedValues[i] = "'" + v + "'"
	}
	valuesStr := ""
	for i, qv := range quotedValues {
		if i > 0 {
			valuesStr += ", "
		}
		valuesStr += qv
	}

	// Build TypeScript union type (e.g., "'admin' | 'user' | 'guest'")
	typeUnion := ""
	for i, qv := range quotedValues {
		if i > 0 {
			typeUnion += " | "
		}
		typeUnion += qv
	}

	return &EnumData{
		Name:       enumName,
		VarName:    varName,
		Values:     enum.Values,
		ValuesStr:  valuesStr,
		TypeUnion:  typeUnion,
		SchemaName: enum.Schema,
	}
}

// NewTableData creates TableData from a models.Table
func NewTableData(table *models.Table, tm *TypeMapper) *TableData {
	tableName := tm.ToCamelCase(table.Name)
	typeName := tm.ToPascalCase(table.Name)

	return &TableData{
		Name:       tableName,
		TableName:  table.Name,
		TypeName:   typeName,
		Columns:    make([]*ColumnData, 0),
		Indexes:    make([]*IndexData, 0),
		Comment:    formatComment(table.Description, table.Comment),
		SchemaName: table.Schema,
	}
}

// AddColumn adds a column to the table data
func (td *TableData) AddColumn(col *ColumnData) {
	td.Columns = append(td.Columns, col)
}

// AddIndex adds an index to the table data
func (td *TableData) AddIndex(idx *IndexData) {
	td.Indexes = append(td.Indexes, idx)
}

// NewColumnData creates ColumnData from a models.Column
func NewColumnData(col *models.Column, table *models.Table, tm *TypeMapper, isEnum bool) *ColumnData {
	fieldName := tm.ToCamelCase(col.Name)
	drizzleChain := tm.BuildColumnChain(col, table, isEnum)

	return &ColumnData{
		Name:         col.Name,
		FieldName:    fieldName,
		DrizzleChain: drizzleChain,
		Comment:      formatComment(col.Description, col.Comment),
	}
}

// NewIndexData creates IndexData from a models.Index
func NewIndexData(index *models.Index, tableVar string, tm *TypeMapper) *IndexData {
	indexName := tm.ToCamelCase(index.Name) + "Idx"

	// Build column references as field names (will be used with destructuring)
	colRefs := make([]string, len(index.Columns))
	for i, colName := range index.Columns {
		// Use just the field name for destructured parameters
		colRefs[i] = tm.ToCamelCase(colName)
	}

	// Build the complete definition
	// Example: index('email_idx').on(email)
	// or: uniqueIndex('unique_email_idx').on(email)
	definition := ""
	if index.Unique {
		definition = "uniqueIndex('" + index.Name + "').on(" + joinStrings(colRefs, ", ") + ")"
	} else {
		definition = "index('" + index.Name + "').on(" + joinStrings(colRefs, ", ") + ")"
	}

	return &IndexData{
		Name:       indexName,
		Columns:    index.Columns,
		IsUnique:   index.Unique,
		Definition: definition,
	}
}

// formatComment combines description and comment into a single comment string
func formatComment(description, comment string) string {
	if description != "" && comment != "" {
		return description + " - " + comment
	}
	if description != "" {
		return description
	}
	return comment
}

// joinStrings joins a slice of strings with a separator
func joinStrings(strs []string, sep string) string {
	result := ""
	for i, s := range strs {
		if i > 0 {
			result += sep
		}
		result += s
	}
	return result
}