PgTidy/pkg/format/body.go

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"
)

// sqlClauseKw: col-0 lines at paren-depth 0 starting with these keywords stay
// as separate logical lines rather than being joined to the previous line.
var sqlClauseKw = map[string]bool{
	"from": true, "where": true, "into": true, "having": true,
	"group": true, "order": true, "returning": true, "set": true,
	"on": true, "join": true, "left": true, "right": true,
	"inner": true, "outer": true, "cross": true, "full": true,
	"union": true, "intersect": true, "except": true,
	"select": true, "with": true,
}

// formatBody applies house-style formatting to a PL/pgSQL dollar-quoted body
// token.
func formatBody(bodyText string, st config.Style) string {
	open, inner, close, ok := splitDollarQuote(bodyText)
	if !ok {
		return bodyText
	}
	return open + formatBodyInner(inner, st) + close
}

// splitDollarQuote splits a dollar-quoted token (e.g. "$$...\n$$" or
// "$S$...$S$") into (open tag, inner text, close tag). The open and close tags
// are the same string; the last occurrence in s is taken as the close tag.
func splitDollarQuote(s string) (open, inner, close string, ok bool) {
	if len(s) == 0 || s[0] != '$' {
		return
	}
	end := strings.Index(s[1:], "$")
	if end < 0 {
		return
	}
	openLen := end + 2
	open = s[:openLen]
	closeStart := strings.LastIndex(s, open)
	if closeStart < openLen {
		return
	}
	inner = s[openLen:closeStart]
	close = s[closeStart:]
	ok = true
	return
}

func formatBodyInner(inner string, st config.Style) string {
	sig, _ := cst.Attach(lexer.Lex(inner))
	nl := st.Newline

	// Find DECLARE at depth 0.
	declareIdx := -1
	for i, t := range sig {
		if t.Tok.Kind == lexer.Ident && lowerASCII(t.Tok.Text) == "declare" {
			declareIdx = i
			break
		}
	}
	if declareIdx < 0 {
		return inner
	}

	// Find BEGIN at depth 0 after DECLARE.
	beginIdx := -1
	depth := 0
	for i := declareIdx + 1; i < len(sig); i++ {
		switch sig[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 && sig[i].Tok.Kind == lexer.Ident && lowerASCII(sig[i].Tok.Text) == "begin" {
			beginIdx = i
			break
		}
	}
	if beginIdx < 0 {
		return inner
	}

	var b strings.Builder

	// Emit verbatim up to and including DECLARE (keyword-cased).
	for i := 0; i <= declareIdx; i++ {
		t := sig[i]
		for _, tr := range t.Lead {
			b.WriteString(tr.Text)
		}
		if i == declareIdx {
			b.WriteString(applyCase(t.Tok.Text, st.KeywordCase))
		} else {
			b.WriteString(t.Tok.Text)
		}
	}
	b.WriteString(nl)

	// Format each variable declaration in the DECLARE section.
	formatDeclareVars(&b, sig[declareIdx+1:beginIdx], st)

	// Format the BEGIN…END block.
	b.WriteString(formatBodyStatements(inner[sig[beginIdx].Tok.Off:], st))

	return b.String()
}

// formatDeclareVars writes each variable declaration as a single indented line.
// Comments in the leading trivia of a declaration's first token are preserved
// on their own lines. If a declaration carries mid-body comments it is emitted
// verbatim to avoid losing them.
func formatDeclareVars(b *strings.Builder, toks []cst.Tok, st config.Style) {
	nl := st.Newline
	indent := st.Indent
	depth := 0
	var cur []cst.Tok
	var preComments []string

	emit := func() {
		if len(cur) == 0 {
			return
		}
		for _, c := range preComments {
			b.WriteString(indent)
			b.WriteString(c)
			b.WriteString(nl)
		}
		preComments = nil

		// Graceful degradation: mid-declaration comments stay verbatim.
		if anyComment(cur[1:]) {
			b.WriteString(indent)
			b.WriteString(verbatimSpan(cur))
			b.WriteString(nl)
			cur = nil
			return
		}

		body := cur
		hasSemi := len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon
		if hasSemi {
			body = body[:len(body)-1]
		}
		b.WriteString(indent)
		for i, t := range body {
			if i > 0 && needSpace(body[i-1].Tok, t.Tok) {
				b.WriteByte(' ')
			}
			b.WriteString(caseText(t.Tok, st))
		}
		if hasSemi {
			b.WriteString(";")
		}
		b.WriteString(nl)
		cur = nil
	}

	for _, t := range toks {
		if len(cur) == 0 {
			for _, tr := range t.Lead {
				if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
					preComments = append(preComments, strings.TrimRight(tr.Text, " \t"))
				}
			}
		}

		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}

		cur = append(cur, t)

		if t.Tok.Kind == lexer.Semicolon && depth == 0 {
			emit()
		}
	}
	emit()
}

// bline is one logical line within an accumulated statement.
type bline struct {
	text   string // content without leading whitespace
	indent string // original leading whitespace
}

// formatBodyStatements formats the BEGIN…END block of a PL/pgSQL body.
// text must begin at the 'B' of BEGIN and include the final END (with optional ;).
//
// Algorithm:
//  1. Physical lines are collected into logical lines, joining col-0 broken
//     continuations (the "split-line join" from the spec).
//  2. Block depth is tracked via PL/pgSQL structural keywords so each statement
//     is indented to depth × st.Indent.
//  3. After EXCEPTION the formatter switches to verbatim-indent mode (original
//     leading whitespace is preserved) to avoid conflicts between styles that
//     put WHEN at col-0 vs indented.
//  4. Blank-line counts from the original are preserved.
func formatBodyStatements(text string, st config.Style) string {
	nl := st.Newline
	normalised := strings.ReplaceAll(text, "\r\n", "\n")
	rawLines := strings.Split(normalised, "\n")

	var (
		result       strings.Builder
		stmt         []bline
		parenDepth   int
		blockDepth   int // 0=col-0 (BEGIN/END/EXCEPTION), 1=body, 2=nested…
		inException  bool
		pendingBlanks int
		depthInc     bool // increment blockDepth after next flush
	)

	flush := func() {
		if len(stmt) == 0 {
			return
		}
		for i := 0; i < pendingBlanks; i++ {
			result.WriteString(nl)
		}
		pendingBlanks = 0

		fw := lowerASCII(firstBodyKeyword(stmt[0].text))

		if inException {
			// Verbatim-indent mode: preserve original leading whitespace.
			for _, ll := range stmt {
				result.WriteString(ll.indent)
				result.WriteString(ll.text)
				result.WriteString(nl)
			}
			stmt = nil
			if depthInc {
				blockDepth++
				depthInc = false
			}
			return
		}

		effectiveDepth := blockDepth
		switch fw {
		case "end":
			blockDepth--
			if blockDepth < 0 {
				blockDepth = 0
			}
			effectiveDepth = blockDepth
		case "else", "elsif", "elseif":
			// Emit at one level up; blockDepth unchanged so the body continues
			// at the same depth (THEN will re-apply depthInc for elsif).
			if blockDepth > 0 {
				effectiveDepth = blockDepth - 1
			}
		case "exception":
			inException = true
			effectiveDepth = 0
		}

		baseIndent := strings.Repeat(st.Indent, effectiveDepth)

		for i, ll := range stmt {
			if i == 0 || ll.indent == "" {
				result.WriteString(baseIndent)
			} else {
				result.WriteString(ll.indent)
			}
			result.WriteString(ll.text)
			result.WriteString(nl)
		}

		if depthInc {
			blockDepth++
			depthInc = false
		}
		stmt = nil
	}

	for _, rawLine := range rawLines {
		line := strings.TrimRight(rawLine, "\r")
		indent := leadingWhitespace(line)
		stripped := line[len(indent):]

		if stripped == "" {
			flush()
			pendingBlanks++
			continue
		}

		fw := lowerASCII(firstBodyKeyword(stripped))
		isColZero := indent == ""
		// Only join when we are at paren-depth 0; content inside parens (e.g.
		// inside a CTE subquery) should not be merged across lines.
		joinToPrev := isColZero && parenDepth == 0 && len(stmt) > 0 && !sqlClauseKw[fw]

		if joinToPrev {
			last := &stmt[len(stmt)-1]
			last.text = strings.TrimRight(last.text, " \t") + " " + stripped
		} else {
			stmt = append(stmt, bline{text: stripped, indent: indent})
		}

		// Scan tokens to track paren depth and detect flush triggers.
		var lastD0Kw string
		for _, tok := range lexer.Lex(stripped) {
			if tok.IsTrivia() || tok.Kind == lexer.EOF {
				continue
			}
			switch tok.Kind {
			case lexer.LParen, lexer.LBracket:
				parenDepth++
			case lexer.RParen, lexer.RBracket:
				if parenDepth > 0 {
					parenDepth--
				}
			case lexer.Semicolon:
				if parenDepth == 0 {
					flush()
				}
			}
			if parenDepth == 0 && tok.Kind == lexer.Ident {
				lastD0Kw = lowerASCII(tok.Text)
			}
		}

		// Structural keywords at the end of a line (paren depth 0) trigger a
		// flush and possibly a block-depth change.
		if parenDepth == 0 && len(stmt) > 0 {
			switch lastD0Kw {
			case "then", "loop", "begin":
				// ELSIF/ELSEIF headers end with THEN but must NOT increment depth
				// (blockDepth is already at the right level for the body).
				fw0 := lowerASCII(firstBodyKeyword(stmt[0].text))
				if fw0 != "elsif" && fw0 != "elseif" {
					depthInc = true
				}
				flush()
			case "else", "exception":
				flush()
			}
		}
	}

	flush()
	return result.String()
}

// firstBodyKeyword returns the text of the first identifier token in s
// (lowercased), or "" if the first significant token is not an identifier.
func firstBodyKeyword(s string) string {
	for _, tok := range lexer.Lex(s) {
		if tok.IsTrivia() || tok.Kind == lexer.EOF {
			continue
		}
		if tok.Kind == lexer.Ident {
			return tok.Text
		}
		return "" // first significant token is non-ident
	}
	return ""
}

// leadingWhitespace returns the leading space/tab prefix of s.
func leadingWhitespace(s string) string {
	i := 0
	for i < len(s) && (s[i] == ' ' || s[i] == '\t') {
		i++
	}
	return s[:i]
}