PgTidy/pkg/format/format.go

// Package format renders a cst.File back to source text in the configured
// house style.
//
// Scope (V1): CREATE FUNCTION/PROCEDURE headers are laid out to house style
// (one parameter per line with leading commas, each option clause on its own
// line, AS/$$ on their own lines). The PL/pgSQL body and any statement the
// parser left as cst.Raw are emitted verbatim — this upholds the safety
// invariants while body formatting is built out (task #4).
//
// Guarantees: formatting changes only trivia/layout (never literal/identifier
// semantics), and is idempotent.
package format

import (
	"strings"

	"git.warky.dev/wdevs/pgtidy/pkg/config"
	"git.warky.dev/wdevs/pgtidy/pkg/cst"
	"git.warky.dev/wdevs/pgtidy/pkg/lexer"
)

// File formats a parsed file with the given style.
func File(f *cst.File, st config.Style) string {
	p := &printer{st: st}
	for i, item := range f.Items {
		toks := cst.Tokens(item)
		if len(toks) == 0 {
			continue
		}
		lead := toks[0].Lead
		if i > 0 {
			p.nl()
			if hasBlankLine(lead) {
				p.nl()
			}
		}
		p.leadingComments(lead)
		p.writeItem(item)
	}
	p.trailingComments(f.Trailing)
	return ensureTrailingNewline(p.b.String(), st.Newline)
}

type printer struct {
	st config.Style
	b  strings.Builder
}

func (p *printer) nl() { p.b.WriteString(p.st.Newline) }

func (p *printer) writeItem(n cst.Node) {
	switch v := n.(type) {
	case *cst.CreateFunction:
		p.writeCreateFunction(v)
	case *cst.Raw:
		switch {
		case isDMLStart(v.Toks):
			p.b.WriteString(formatDML(v.Toks, p.st))
		case isDoBlock(v.Toks):
			p.b.WriteString(formatDoBlock(v.Toks, p.st))
		default:
			p.b.WriteString(verbatimSpan(v.Toks))
		}
	default:
		p.b.WriteString(verbatimSpan(cst.Tokens(n)))
	}
}

// isDoBlock reports whether toks is a DO $$ ... $$ statement.
func isDoBlock(toks []cst.Tok) bool {
	for _, t := range toks {
		if t.Tok.Kind == lexer.Ident {
			return lowerASCII(t.Tok.Text) == "do"
		}
		if !t.Tok.IsTrivia() {
			return false
		}
	}
	return false
}

// formatDoBlock formats a DO $$ ... $$ block by applying formatBody to the
// dollar-quoted string and emitting DO + newline + formatted body.
func formatDoBlock(toks []cst.Tok, st config.Style) string {
	// Find the DO keyword, the dollar-string body, and the optional semicolon.
	var doTok, bodyTok *cst.Tok
	hasSemi := false
	for i := range toks {
		t := &toks[i]
		if t.Tok.IsTrivia() || t.Tok.Kind == lexer.EOF {
			continue
		}
		low := lowerASCII(t.Tok.Text)
		if t.Tok.Kind == lexer.Ident && low == "do" && doTok == nil {
			doTok = t
			continue
		}
		if doTok != nil && t.Tok.Kind == lexer.DollarString && bodyTok == nil {
			bodyTok = t
			continue
		}
		if t.Tok.Kind == lexer.Semicolon {
			hasSemi = true
		}
	}
	if doTok == nil || bodyTok == nil {
		return verbatimSpan(toks)
	}

	nl := st.Newline
	var b strings.Builder
	b.WriteString(applyCase(doTok.Tok.Text, st.KeywordCase))
	b.WriteString(nl)
	b.WriteString(formatBody(bodyTok.Tok.Text, st))
	if hasSemi {
		b.WriteString(";")
	}
	return b.String()
}

func (p *printer) writeCreateFunction(cf *cst.CreateFunction) {
	// Safety: if the header carries comments we cannot confidently relocate,
	// emit the whole statement verbatim rather than risk dropping them.
	// We still format the body dollar-string independently since it is self-contained.
	if headerHasComments(cf) {
		p.b.WriteString(verbatimSpanFormatBody(cst.Tokens(cf), cf.Body, p.st))
		return
	}

	head := p.inline(cf.Head)
	name := p.inline(cf.Name)
	p.b.WriteString(head)
	if name != "" {
		p.b.WriteString(" ")
		p.b.WriteString(name)
	}
	p.b.WriteString("(")

	// Build formatted param texts first so we can measure widths.
	paramTexts := make([]string, len(cf.Params))
	for i, param := range cf.Params {
		paramTexts[i] = p.inline(param.Toks)
	}

	// align_param_types: pad param names so type columns align.
	if p.st.AlignParamTypes && len(cf.Params) > 1 {
		paramTexts = alignParamTypes(paramTexts)
	}

	first := p.st.Indent + " " // align item text one column past the comma
	cont := p.st.Indent
	for i, text := range paramTexts {
		param := cf.Params[i]
		p.nl()
		if i == 0 || p.st.Commas != config.CommaLeading {
			p.b.WriteString(first)
			p.b.WriteString(text)
			if p.st.Commas == config.CommaTrailing && i < len(cf.Params)-1 {
				p.b.WriteString(",")
			}
		} else {
			p.b.WriteString(cont)
			p.b.WriteString(",")
			p.b.WriteString(text)
		}
		// Emit trailing inline comment from the separator (e.g. --description after param).
		if param.Sep != nil {
			for _, tr := range param.Sep.Lead {
				if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
					p.b.WriteByte(' ')
					p.b.WriteString(strings.TrimRight(tr.Text, " \t"))
					break
				}
			}
		}
	}
	p.nl()
	p.b.WriteString(")")

	for _, clause := range cf.Options {
		p.nl()
		p.b.WriteString(p.inline(clause))
	}
	if cf.As != nil {
		// routine_as_wrap: when false, AS stays on the same line as the last option.
		if p.st.RoutineAsWrap {
			p.nl()
		} else {
			p.b.WriteByte(' ')
		}
		p.b.WriteString(p.inline([]cst.Tok{{Tok: cf.As.Tok}}))
	}
	if cf.Body != nil {
		p.nl()
		p.b.WriteString(formatBody(cf.Body.Tok.Text, p.st))
	}
	for _, clause := range cf.Tail {
		p.nl()
		p.b.WriteString(p.inline(clause))
	}
	if cf.Semi != nil {
		p.b.WriteString(";")
	}
}

// inline renders a run of tokens on one line, applying spacing and casing.
// If the run contains comment trivia it is emitted verbatim to avoid losing
// or misplacing the comments.
func (p *printer) inline(toks []cst.Tok) string {
	if len(toks) == 0 {
		return ""
	}
	// A comment on the first token's lead is the span's leading/separation
	// comment, handled by the caller and never emitted here, so it does not
	// force verbatim. Internal comments (on later tokens) do.
	if anyComment(toks[1:]) {
		return verbatimSpan(toks)
	}
	var b strings.Builder
	for i, t := range toks {
		if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
			b.WriteByte(' ')
		}
		var prev lexer.Token
		if i > 0 {
			prev = toks[i-1].Tok
		}
		nextIsLParen := i+1 < len(toks) && toks[i+1].Tok.Kind == lexer.LParen
		b.WriteString(caseTextCtx(t.Tok, prev, nextIsLParen, p.st))
	}
	return b.String()
}

func (p *printer) leadingComments(lead cst.Trivia) {
	for _, tr := range lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			text := strings.TrimRight(strings.ReplaceAll(tr.Text, "\r", ""), " \t")
			p.b.WriteString(text)
			p.nl()
		}
	}
}

func (p *printer) trailingComments(lead cst.Trivia) {
	cs := commentsOf(lead)
	for _, c := range cs {
		p.nl()
		p.b.WriteString(strings.TrimRight(strings.ReplaceAll(c.Text, "\r", ""), " \t"))
	}
}

// --- spacing & casing ---

// tightOps are operators printed without surrounding spaces.
var tightOps = map[string]bool{"::": true, ":": true, "->": true, "->>": true}

// parenKws are keywords that always take a space before '(' because they
// introduce a subquery or a bracketed clause, not a function-call argument list.
var parenKws = map[string]bool{
	"as": true, "exists": true, "in": true, "not": true,
	"between": true, "like": true, "ilike": true, "similar": true,
}

func needSpace(a, b lexer.Token) bool {
	// No space after.
	switch a.Kind {
	case lexer.LParen, lexer.LBracket, lexer.Dot:
		return false
	case lexer.Operator:
		if tightOps[a.Text] {
			return false
		}
	}
	// No space before.
	switch b.Kind {
	case lexer.RParen, lexer.RBracket, lexer.Comma, lexer.Semicolon, lexer.Dot:
		return false
	case lexer.LParen:
		switch a.Kind {
		case lexer.Ident, lexer.QuotedIdent, lexer.RParen, lexer.RBracket, lexer.Param:
			// Keywords like AS/EXISTS/IN introduce subqueries or clause groups,
			// not function calls — always separate with a space.
			if a.Kind == lexer.Ident && parenKws[lowerASCII(a.Text)] {
				return true
			}
			return false // function call / type modifier
		}
	case lexer.LBracket:
		switch a.Kind {
		case lexer.Ident, lexer.QuotedIdent, lexer.RParen, lexer.RBracket:
			return false // array subscript / array type modifier
		}
	case lexer.Operator:
		if tightOps[b.Text] {
			return false
		}
	}
	return true
}

func caseText(t lexer.Token, st config.Style) string {
	return caseTextCtx(t, lexer.Token{}, false, st)
}

// caseTextCtx applies casing with context: prev is the preceding significant
// token, nextIsLParen indicates the next significant token is '('.
func caseTextCtx(t lexer.Token, prev lexer.Token, nextIsLParen bool, st config.Style) string {
	if t.Kind != lexer.Ident {
		return t.Text
	}
	low := lowerASCII(t.Text)
	switch {
	case isTypeName(low):
		return applyCase(t.Text, st.TypeCase)
	case isKeyword(low):
		return applyCase(t.Text, st.KeywordCase)
	case nextIsLParen && isBuiltinFunc(low):
		return applyCase(t.Text, st.BuiltinCase)
	case prev.Kind == lexer.Ident && lowerASCII(prev.Text) == "as":
		return applyCase(t.Text, st.AliasCase)
	default:
		return applyCase(t.Text, st.IdentCase)
	}
}

func applyCase(s string, c config.Case) string {
	switch c {
	case config.CaseUpper:
		return strings.ToUpper(s)
	case config.CaseLower:
		return strings.ToLower(s)
	default:
		return s
	}
}

// --- helpers ---

// headerHasComments reports whether the function header carries comment trivia
// the formatter cannot confidently relocate. Param separator (Sep) comments
// are excluded — those are trailing inline comments on param lines that the
// formatter emits explicitly after each param text. The first token's leading
// trivia and the body token are also excluded.
func headerHasComments(cf *cst.CreateFunction) bool {
	// Build a set of Sep token offsets so we can skip them.
	sepOffsets := make(map[int]bool, len(cf.Params))
	for _, p := range cf.Params {
		if p.Sep != nil {
			sepOffsets[p.Sep.Tok.Off] = true
		}
	}
	all := cst.Tokens(cf)
	for i, t := range all {
		if i == 0 || t.Tok.Kind == lexer.Semicolon {
			continue
		}
		if sepOffsets[t.Tok.Off] {
			continue // Sep comments handled separately
		}
		if hasComment(t) {
			return true
		}
	}
	return false
}

func anyComment(toks []cst.Tok) bool {
	for _, t := range toks {
		if hasComment(t) {
			return true
		}
	}
	return false
}

func hasComment(t cst.Tok) bool {
	for _, tr := range t.Lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			return true
		}
	}
	return false
}

func commentsOf(lead cst.Trivia) []lexer.Token {
	var out []lexer.Token
	for _, tr := range lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			out = append(out, tr)
		}
	}
	return out
}

// verbatimSpan emits tokens exactly as in source, excluding the leading trivia
// of the first token (separation is controlled by the caller).
func verbatimSpan(toks []cst.Tok) string {
	var b strings.Builder
	for i, t := range toks {
		if i > 0 {
			for _, tr := range t.Lead {
				b.WriteString(tr.Text)
			}
		}
		b.WriteString(t.Tok.Text)
	}
	return b.String()
}

// verbatimSpanFormatBody emits toks verbatim but replaces bodyTok's text with
// formatBody output. Used when the function header has comments we cannot
// safely relocate but the body can still be independently formatted.
// If bodyTok is nil the function is identical to verbatimSpan.
func verbatimSpanFormatBody(toks []cst.Tok, bodyTok *cst.Tok, st config.Style) string {
	var b strings.Builder
	for i, t := range toks {
		if i > 0 {
			for _, tr := range t.Lead {
				b.WriteString(tr.Text)
			}
		}
		if bodyTok != nil && t.Tok.Kind == lexer.DollarString && t.Tok.Off == bodyTok.Tok.Off {
			b.WriteString(formatBody(t.Tok.Text, st))
		} else {
			b.WriteString(t.Tok.Text)
		}
	}
	return b.String()
}

// hasBlankLine reports whether leading whitespace trivia contains a blank line
// (two or more newlines), indicating the author wanted statements separated.
func hasBlankLine(lead cst.Trivia) bool {
	n := 0
	for _, tr := range lead {
		if tr.Kind == lexer.Whitespace {
			n += strings.Count(tr.Text, "\n")
		}
	}
	return n >= 2
}

func ensureTrailingNewline(s, nl string) string {
	s = strings.TrimRight(s, "\n")
	if s == "" {
		return s
	}
	return s + nl
}

func lowerASCII(s string) string {
	b := []byte(s)
	for i, c := range b {
		if c >= 'A' && c <= 'Z' {
			b[i] = c + 32
		}
	}
	return string(b)
}

// alignParamTypes pads param names so the type column aligns across all params.
// Expected format per param: "[mode] name type [DEFAULT expr]".
// Mode keywords (IN/OUT/INOUT/VARIADIC) are detected and skipped.
// Params without a type are passed through unchanged.
func alignParamTypes(params []string) []string {
	type pp struct{ mode, name, rest string }
	parsed := make([]pp, len(params))
	maxNameW := 0

	modeKws := map[string]bool{"in": true, "out": true, "inout": true, "variadic": true}

	for i, s := range params {
		fields := strings.Fields(s)
		if len(fields) < 2 {
			parsed[i].rest = s
			continue
		}
		nameIdx := 0
		if modeKws[lowerASCII(fields[0])] {
			nameIdx = 1
		}
		if nameIdx >= len(fields) || nameIdx+1 >= len(fields) {
			// No type field — keep verbatim.
			parsed[i].rest = s
			continue
		}
		if nameIdx > 0 {
			parsed[i].mode = fields[0]
		}
		parsed[i].name = fields[nameIdx]
		parsed[i].rest = strings.Join(fields[nameIdx+1:], " ")
		if len(parsed[i].name) > maxNameW {
			maxNameW = len(parsed[i].name)
		}
	}

	if maxNameW == 0 {
		return params
	}

	out := make([]string, len(params))
	for i, p := range parsed {
		if p.name == "" {
			out[i] = params[i]
			continue
		}
		var b strings.Builder
		if p.mode != "" {
			b.WriteString(p.mode)
			b.WriteByte(' ')
		}
		b.WriteString(p.name)
		// Pad name to (maxNameW+1) so the type column starts at a consistent offset.
		pad := maxNameW + 1 - len(p.name)
		for k := 0; k < pad; k++ {
			b.WriteByte(' ')
		}
		b.WriteString(p.rest)
		out[i] = b.String()
	}
	return out
}