PgTidy/pkg/format/body.go

package format

import (
	"git.warky.dev/wdevs/pgtidy/pkg/parser"
	"regexp"
	"sort"
	"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).
	// Normalize CRLF in trivia so the output always uses st.Newline.
	for i := 0; i <= declareIdx; i++ {
		t := sig[i]
		for _, tr := range t.Lead {
			b.WriteString(strings.ReplaceAll(tr.Text, "\r\n", nl))
		}
		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)

	// Comments between the last declaration and BEGIN live in BEGIN's leading
	// trivia; keep them, one per line, at declaration indent.
	for _, tr := range sig[beginIdx].Lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			b.WriteString(st.Indent)
			b.WriteString(strings.TrimRight(strings.ReplaceAll(tr.Text, "\r", ""), " \t"))
			b.WriteString(nl)
		}
	}

	// 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 decls [][]cst.Tok
	var cur []cst.Tok
	var preCommentSets [][]string
	var curPreComments []string

	collect := func() {
		if len(cur) == 0 {
			return
		}
		decls = append(decls, cur)
		preCommentSets = append(preCommentSets, curPreComments)
		cur = nil
		curPreComments = nil
	}

	for _, t := range toks {
		if len(cur) == 0 {
			for _, tr := range t.Lead {
				if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
					text := strings.TrimRight(strings.ReplaceAll(tr.Text, "\r", ""), " \t")
					curPreComments = append(curPreComments, text)
				}
			}
		}
		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 {
			collect()
		}
	}
	collect()

	// Compute alignment widths when requested.
	var nameColW, typeColW int
	if st.PlpgsqlDeclareAlignType || st.PlpgsqlDeclareAlignEq {
		for _, decl := range decls {
			if anyComment(decl[1:]) {
				continue
			}
			body := decl
			if len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon {
				body = body[:len(body)-1]
			}
			nw, tw := declareNameTypeWidth(body, st)
			if nw > nameColW {
				nameColW = nw
			}
			// typeColW drives the '='/':='/DEFAULT column (align_eq only), so
			// only declarations that actually carry an assignment participate —
			// a bare "name type;" must not widen it.
			if tw > typeColW && declHasAssignment(body) {
				typeColW = tw
			}
		}
	}

	for i, cur := range decls {
		for _, c := range preCommentSets[i] {
			b.WriteString(indent)
			b.WriteString(c)
			b.WriteString(nl)
		}
		// Graceful degradation: mid-declaration comments stay verbatim.
		if anyComment(cur[1:]) {
			b.WriteString(indent)
			b.WriteString(verbatimSpan(cur))
			b.WriteString(nl)
			continue
		}

		body := cur
		hasSemi := len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon
		if hasSemi {
			body = body[:len(body)-1]
		}
		b.WriteString(indent)
		if (st.PlpgsqlDeclareAlignType || st.PlpgsqlDeclareAlignEq) && nameColW > 0 {
			writeDeclareAligned(b, body, st, nameColW, typeColW)
		} else {
			for j, t := range body {
				if j > 0 && needSpace(body[j-1].Tok, t.Tok) && !isPctTypeBoundary(body, j) {
					b.WriteByte(' ')
				}
				b.WriteString(caseText(t.Tok, st))
			}
		}
		if hasSemi {
			b.WriteString(";")
		}
		b.WriteString(nl)
	}
}

// declHasAssignment reports whether a DECLARE variable body (name type … ) has
// a default assignment ( := / = / DEFAULT ) at paren depth 0.
func declHasAssignment(body []cst.Tok) bool {
	depth := 0
	for _, t := range body {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth != 0 {
			continue
		}
		if t.Tok.Kind == lexer.Operator && (t.Tok.Text == ":=" || t.Tok.Text == "=") {
			return true
		}
		if t.Tok.Kind == lexer.Ident && lowerASCII(t.Tok.Text) == "default" {
			return true
		}
	}
	return false
}

// declareNameTypeWidth returns the rendered width of the name and type portions
// of a DECLARE variable declaration (without the default assignment).
// Format is: [name type [:= default]] or [name type [DEFAULT default]].
func declareNameTypeWidth(body []cst.Tok, st config.Style) (nameW, typeW int) {
	if len(body) < 2 {
		return 0, 0
	}
	// name is always the first token.
	name := caseText(body[0].Tok, st)
	nameW = len(name)

	// type runs from body[1] until we hit := / DEFAULT / = at depth 0.
	var typeTokens []cst.Tok
	depth := 0
	for _, t := range body[1:] {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 {
			low := lowerASCII(t.Tok.Text)
			if t.Tok.Kind == lexer.Operator && (t.Tok.Text == ":=" || t.Tok.Text == "=") {
				break
			}
			if t.Tok.Kind == lexer.Ident && low == "default" {
				break
			}
		}
		typeTokens = append(typeTokens, t)
	}
	var tb strings.Builder
	for j, t := range typeTokens {
		if j > 0 && needSpace(typeTokens[j-1].Tok, t.Tok) && !isPctTypeBoundary(typeTokens, j) {
			tb.WriteByte(' ')
		}
		tb.WriteString(caseText(t.Tok, st))
	}
	typeW = len(tb.String())
	return nameW, typeW
}

// writeDeclareAligned writes a single DECLARE variable with aligned columns.
func writeDeclareAligned(b *strings.Builder, body []cst.Tok, st config.Style, nameColW, typeColW int) {
	if len(body) == 0 {
		return
	}
	name := caseText(body[0].Tok, st)
	b.WriteString(name)
	if len(body) == 1 {
		return
	}

	// Pad name to nameColW if align_type is requested.
	if st.PlpgsqlDeclareAlignType {
		pad := nameColW - len(name)
		for k := 0; k < pad; k++ {
			b.WriteByte(' ')
		}
	}

	// Collect type tokens.
	var typeTokens, restTokens []cst.Tok
	depth := 0
	pastType := false
	for _, t := range body[1:] {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if !pastType && depth == 0 {
			low := lowerASCII(t.Tok.Text)
			if (t.Tok.Kind == lexer.Operator && (t.Tok.Text == ":=" || t.Tok.Text == "=")) ||
				(t.Tok.Kind == lexer.Ident && low == "default") {
				pastType = true
				restTokens = append(restTokens, t)
				continue
			}
		}
		if pastType {
			restTokens = append(restTokens, t)
		} else {
			typeTokens = append(typeTokens, t)
		}
	}

	var typeStr strings.Builder
	for j, t := range typeTokens {
		if j > 0 && needSpace(typeTokens[j-1].Tok, t.Tok) && !isPctTypeBoundary(typeTokens, j) {
			typeStr.WriteByte(' ')
		}
		typeStr.WriteString(caseText(t.Tok, st))
	}
	typeRendered := typeStr.String()

	b.WriteByte(' ')
	b.WriteString(typeRendered)

	if len(restTokens) > 0 {
		// Pad type to typeColW if align_eq is requested.
		if st.PlpgsqlDeclareAlignEq {
			pad := typeColW - len(typeRendered)
			for k := 0; k < pad; k++ {
				b.WriteByte(' ')
			}
		}
		for j, t := range restTokens {
			prev := restTokens[0].Tok
			if j > 0 {
				prev = restTokens[j-1].Tok
			}
			if j == 0 || needSpace(prev, t.Tok) {
				b.WriteByte(' ')
			}
			b.WriteString(caseText(t.Tok, st))
		}
	}
}

// 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 (capped by PlpgsqlMaxBlankLines).
func formatBodyStatements(text string, st config.Style) string {
	nl := st.Newline
	text = formatEmbeddedLiterals(text, st)
	normalised := strings.ReplaceAll(text, "\r\n", "\n")
	rawLines := strings.Split(normalised, "\n")

	// Mark the continuation lines of every multi-line token whose interior is
	// not code: /* … */ block comments and every kind of string literal,
	// including dollar-quoted ones (a dollar quote is just another way of
	// quoting a string, whatever its tag, and its content may be any language).
	// Those lines are token content, not statements: they must be carried
	// verbatim with the token's opening line, never split off and reindented.
	// On the line where such a token ends, tail holds the code that follows it
	// so paren depth and statement terminators are still tracked.
	inVerbatim := make([]bool, len(rawLines))
	tail := make([]string, len(rawLines))
	tailCol := make([]int, len(rawLines)) // column where tail[i] starts
	cut := make([]int, len(rawLines))     // column where a multi-line token opens on this line, or -1
	for i := range cut {
		cut[i] = -1
	}
	lineStart := make([]int, len(rawLines))
	for i, off := 0, 0; i < len(rawLines); i++ {
		lineStart[i] = off
		off += len(rawLines[i]) + 1
	}
	lineOf := func(off int) int {
		return sort.Search(len(lineStart), func(i int) bool { return lineStart[i] > off }) - 1
	}
	for _, t := range lexer.Lex(normalised) {
		switch t.Kind {
		case lexer.BlockComment, lexer.String, lexer.EscapeString, lexer.BitString, lexer.DollarString:
		default:
			continue
		}
		n := strings.Count(t.Text, "\n")
		if n == 0 || !literalTerminated(t) {
			continue
		}
		start := lineOf(t.Off)
		if cut[start] < 0 {
			cut[start] = t.Off - lineStart[start]
		}
		for k := 1; k <= n && start+k < len(inVerbatim); k++ {
			inVerbatim[start+k] = true
		}
		if start+n < len(tail) {
			rest := normalised[t.Off+len(t.Text):]
			if i := strings.IndexByte(rest, '\n'); i >= 0 {
				rest = rest[:i]
			}
			tail[start+n] = rest
			tailCol[start+n] = t.Off + len(t.Text) - lineStart[start+n]
		}
	}

	maxBlanks := st.PlpgsqlMaxBlankLines
	if maxBlanks < 0 {
		maxBlanks = 0
	}

	var (
		result        strings.Builder
		stmt          []bline
		parenDepth    int
		blockDepth    int // 0=col-0 (BEGIN/END/EXCEPTION), 1=body, 2=nested…
		caseDepth     int // depth of open CASE…END expressions (WHEN…THEN is not a block opener)
		inException   bool
		pendingBlanks int
		depthInc      bool // increment blockDepth after next flush
	)

	flush := func() {
		if len(stmt) == 0 {
			return
		}
		blanks := pendingBlanks
		if blanks > maxBlanks {
			blanks = maxBlanks
		}
		for i := 0; i < blanks; 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":
			if blockDepth > 0 {
				effectiveDepth = blockDepth - 1
			}
		case "exception":
			inException = true
			// EXCEPTION belongs to its nearest enclosing BEGIN, so align it one
			// level in from the current block body (col 0 for the outermost).
			effectiveDepth = blockDepth - 1
			if effectiveDepth < 0 {
				effectiveDepth = 0
			}
		}

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

		// plpgsql_if_then_newline: when false, THEN stays on the same line as
		// the condition. When true (default) it's already on its own logical line.
		stmtLines := stmt
		if !st.PlpgsqlIfThenNewline && fw == "if" {
			stmtLines = joinThenToCondition(stmt)
		}

		formattedLines := formatBodyStmtLines(stmtLines, baseIndent, st)
		for _, line := range formattedLines {
			result.WriteString(line)
			result.WriteString(nl)
		}

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

	scanLine := func(text string, openLit bool) {
		var lastD0Kw string
		for _, tok := range lexer.Lex(text) {
			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 {
					// plpgsql_loop_collapse: fold empty FOR … LOOP END LOOP; to one line.
					if st.PlpgsqlLoopCollapse && len(stmt) > 0 {
						collapsed, ok := tryCollapseLoop(stmt, st)
						if ok {
							stmt = []bline{{text: collapsed, indent: ""}}
						}
					}
					flush()
				}
			}
			if parenDepth == 0 && tok.Kind == lexer.Ident {
				lastD0Kw = lowerASCII(tok.Text)
				switch lastD0Kw {
				case "case":
					caseDepth++
				case "end":
					if caseDepth > 0 {
						caseDepth--
					}
				}
			} else if parenDepth == 0 {
				// A non-keyword token (string, operator, …) ends the run of
				// keywords, so a line like `raise exception 'x'` does not end
				// in the keyword EXCEPTION.
				lastD0Kw = ""
			}
		}

		if openLit {
			lastD0Kw = "" // the line ends inside a literal, not on a keyword
		}
		if parenDepth == 0 && len(stmt) > 0 {
			switch lastD0Kw {
			case "then":
				// A THEN ending a CASE…WHEN branch is not a PL/pgSQL block
				// opener; only one matching END closes the whole CASE, so
				// treating each WHEN…THEN as a block open would permanently
				// inflate blockDepth.
				if caseDepth == 0 {
					fw0 := lowerASCII(firstBodyKeyword(stmt[0].text))
					if fw0 != "elsif" && fw0 != "elseif" {
						depthInc = true
					}
					flush()
				}
			case "loop", "begin":
				fw0 := lowerASCII(firstBodyKeyword(stmt[0].text))
				if fw0 != "elsif" && fw0 != "elseif" {
					depthInc = true
				}
				flush()
			case "else", "exception":
				if caseDepth > 0 {
					break
				}
				flush()
			}
		}
	}

	for j, rawLine := range rawLines {
		line := strings.TrimRight(rawLine, "\r")

		if inVerbatim[j] {
			// Verbatim continuation of a multi-line comment or literal: glue it
			// to the bline holding the token's opening line.
			if len(stmt) > 0 {
				last := &stmt[len(stmt)-1]
				last.text += "\n" + line
			} else {
				stmt = append(stmt, bline{text: line})
			}
			if tail[j] != "" {
				// A second literal may open later on this same line; scan only
				// the code between the two.
				tt, openLit := tail[j], false
				if c := cut[j] - tailCol[j]; cut[j] >= 0 && c < len(tt) {
					if c < 0 {
						c = 0
					}
					tt, openLit = tt[:c], true
				}
				scanLine(tt, openLit)
			}
			continue
		}

		indent := leadingWhitespace(line)
		stripped := line[len(indent):]

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

		fw := lowerASCII(firstBodyKeyword(stripped))
		isColZero := indent == ""
		joinToPrev := isColZero && parenDepth == 0 && len(stmt) > 0 && !sqlClauseKw[fw]
		// A col-0 comment-only line is its own thing: never glue it onto the
		// previous line — doing so buries a code line's trailing text in a
		// comment and collapses consecutive -- comment lines into one.
		if joinToPrev && len(significantBodyTokens(stripped)) == 0 {
			joinToPrev = false
		}
		// Don't join a col-0 continuation to a comment-only preceding bline:
		// the comment has no structural keyword so `continue ;` at col-0 would
		// disappear into the comment text and be invisible to the lexer.
		if joinToPrev {
			if lowerASCII(firstBodyKeyword(stmt[len(stmt)-1].text)) == "" {
				joinToPrev = false
			}
		}
		// Pre-flush pending comment-only blines before adding a new non-comment
		// non-joined bline. Without this, a comment + `end if;` end up in the
		// same stmt, `fw` comes from the comment (empty string), depth is never
		// decremented, and the formatter diverges on the second pass.
		if !joinToPrev && fw != "" && len(stmt) > 0 {
			allComments := true
			for _, ll := range stmt {
				if lowerASCII(firstBodyKeyword(ll.text)) != "" {
					allComments = false
					break
				}
			}
			if allComments {
				flush()
			}
		}

		// Likewise never join onto a line that ends in a -- comment: the joined
		// code would become part of the comment text.
		if joinToPrev && endsInLineComment(stmt[len(stmt)-1].text) {
			joinToPrev = false
		}

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

		// A multi-line literal that opens on this line is scanned only up to
		// its opening quote; the rest is token content, not code.
		scanText, openLit := stripped, false
		if c := cut[j] - len(indent); cut[j] >= 0 && c < len(scanText) {
			if c < 0 {
				c = 0
			}
			scanText, openLit = scanText[:c], true
		}
		scanLine(scanText, openLit)
	}

	flush()
	return result.String()
}

// formatBodyStmtLines formats one flushed PL/pgSQL statement at its contextual
// base indent. Multi-line UPDATE/DELETE statements inside PL/pgSQL get their
// top-level SET/WHERE/AND/OR clauses realigned under the statement while nested
// subqueries keep their original indentation. Non-DML statements keep
// continuation indentation, except that standalone structural keywords such as
// THEN are aligned with the block opener.
func formatBodyStmtLines(lines []bline, baseIndent string, st config.Style) []string {
	if len(lines) == 0 {
		return nil
	}

	if looksLikeMultiLineBodyDML(lines) {
		return reindentBodyDML(lines, baseIndent, st)
	}

	out := make([]string, 0, len(lines))
	for i, ll := range lines {
		text := ll.text
		indent := baseIndent
		switch {
		case i > 0 && ll.indent == "" && len(significantBodyTokens(text)) == 0:
			// A column-0 comment line trailing a multi-line (commented-out)
			// statement is a continuation the author left flush-left — keep it
			// there rather than re-indenting it to block depth.
			indent = ""
		case i > 0 && ll.indent != "" && !isStandaloneBodyKeyword(ll.text, "then", "else", "elsif", "elseif"):
			indent = ll.indent
		}
		out = append(out, indent+text)
	}
	return out
}

func looksLikeMultiLineBodyDML(lines []bline) bool {
	if len(lines) < 2 {
		return false
	}
	kw := lowerASCII(firstBodyKeyword(lines[0].text))
	return kw == "update" || kw == "delete"
}

func reindentBodyDML(lines []bline, baseIndent string, st config.Style) []string {
	out := make([]string, 0, len(lines)+1)
	afterWhere := false
	parenDepth := 0
	for i, ll := range lines {
		text := strings.TrimRight(ll.text, " 	")
		lineDepth := parenDepth
		kw := lowerASCII(firstBodyKeyword(text))
		if afterWhere && lineDepth == 0 && kw != "and" && kw != "or" {
			out = append(out, baseIndent+st.Indent+st.Indent+strings.TrimSpace(text))
			afterWhere = false
			updateBodyParenDepth(text, &parenDepth)
			continue
		}
		if lineDepth == 0 && (kw == "set" || kw == "where" || kw == "values" || kw == "returning") {
			if kw == "where" {
				whereText := strings.TrimSpace(text)
				fields := strings.Fields(whereText)
				nextKw := ""
				if i+1 < len(lines) {
					nextKw = lowerASCII(firstBodyKeyword(lines[i+1].text))
				}
				if len(fields) > 1 && (nextKw == "and" || nextKw == "or") {
					out = append(out, baseIndent+fields[0])
					out = append(out, baseIndent+st.Indent+st.Indent+strings.TrimSpace(whereText[len(fields[0]):]))
					afterWhere = false
					continue
				}
				afterWhere = len(fields) == 1
			}
			out = append(out, baseIndent+strings.TrimSpace(text))
			continue
		}
		if lineDepth == 0 && (kw == "and" || kw == "or") {
			// House style: AND/OR line up with the WHERE keyword; the first
			// predicate is indented two levels under it.
			out = append(out, baseIndent+strings.TrimSpace(text))
			afterWhere = false
			updateBodyParenDepth(text, &parenDepth)
			continue
		}
		if i == 0 {
			out = append(out, baseIndent+strings.TrimSpace(text))
		} else if ll.indent != "" {
			out = append(out, ll.indent+strings.TrimSpace(text))
		} else {
			out = append(out, baseIndent+strings.TrimSpace(text))
		}
		afterWhere = false
		updateBodyParenDepth(text, &parenDepth)
	}
	return out
}

func updateBodyParenDepth(s string, depth *int) {
	for _, tok := range lexer.Lex(s) {
		switch tok.Kind {
		case lexer.LParen, lexer.LBracket:
			(*depth)++
		case lexer.RParen, lexer.RBracket:
			if *depth > 0 {
				(*depth)--
			}
		}
	}
}

func significantBodyTokens(s string) []cst.Tok {
	var toks []cst.Tok
	for _, tok := range lexer.Lex(s) {
		if tok.IsTrivia() || tok.Kind == lexer.EOF {
			continue
		}
		toks = append(toks, cst.Tok{Tok: tok})
	}
	return toks
}

func isStandaloneBodyKeyword(s string, kws ...string) bool {
	toks := significantBodyTokens(s)
	if len(toks) != 1 || toks[0].Tok.Kind != lexer.Ident {
		return false
	}
	low := lowerASCII(toks[0].Tok.Text)
	for _, kw := range kws {
		if low == kw {
			return true
		}
	}
	return false
}

// joinThenToCondition merges a THEN line (on its own bline) into the preceding
// condition line when plpgsql_if_then_newline is false.
func joinThenToCondition(lines []bline) []bline {
	out := make([]bline, 0, len(lines))
	for i, ll := range lines {
		if i > 0 && strings.EqualFold(strings.TrimSpace(ll.text), "then") {
			out[len(out)-1].text = strings.TrimRight(out[len(out)-1].text, " \t") + " THEN"
		} else {
			out = append(out, ll)
		}
	}
	return out
}

// tryCollapseLoop tries to collapse an empty loop body to one line.
// Detects: FOR … LOOP\n  (empty or only blanks)\nEND LOOP;
// Returns the collapsed line and true on success.
func tryCollapseLoop(lines []bline, st config.Style) (string, bool) {
	if len(lines) < 2 {
		return "", false
	}
	first := strings.TrimSpace(lines[0].text)
	last := strings.TrimSpace(lines[len(lines)-1].text)
	firstLow := lowerASCII(first)
	lastLow := lowerASCII(last)

	// Check last line is END LOOP; or LOOP (for WHILE/FOR empty bodies that end with LOOP).
	if !strings.HasPrefix(lastLow, "end loop") && lastLow != "end loop;" {
		return "", false
	}
	// Check middle lines are all empty.
	for _, mid := range lines[1 : len(lines)-1] {
		if strings.TrimSpace(mid.text) != "" {
			return "", false
		}
	}
	// Check first line ends with LOOP.
	if !strings.HasSuffix(firstLow, "loop") {
		return "", false
	}
	_ = st
	_ = firstLow
	// Collapse to: <header> END LOOP;
	return strings.TrimRight(first, " \t") + " " + strings.ToUpper(last), true
}

// 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]
}

// endsInLineComment reports whether s ends with a -- line comment.
func endsInLineComment(s string) bool {
	toks := lexer.Lex(s)
	for i := len(toks) - 1; i >= 0; i-- {
		switch toks[i].Kind {
		case lexer.EOF, lexer.Whitespace:
			continue
		}
		return toks[i].Kind == lexer.LineComment
	}
	return false
}

// literalTerminated reports whether a string or block-comment token is closed.
// Slices of nested dollar-quoted text can end mid-literal; an unterminated
// token there is a lexing artefact, not real literal content.
func literalTerminated(t lexer.Token) bool {
	if t.Kind == lexer.BlockComment {
		return strings.HasSuffix(t.Text, "*/") && len(t.Text) >= 4
	}
	if t.Kind == lexer.DollarString {
		_, _, _, ok := splitDollarQuote(t.Text)
		return ok
	}
	return len(t.Text) >= 2 && strings.HasSuffix(t.Text, "'")
}

// formatEmbeddedLiterals restyles every dollar-quoted literal in text whose
// content looks like code (see formatEmbedded). Everything else is untouched.
func formatEmbeddedLiterals(text string, st config.Style) string {
	if !strings.Contains(text, "$") {
		return text
	}
	var b strings.Builder
	for _, t := range lexer.Lex(text) {
		if t.Kind == lexer.DollarString {
			b.WriteString(formatEmbedded(t.Text, st))
		} else {
			b.WriteString(t.Text)
		}
	}
	return b.String()
}

// formatEmbedded formats a dollar-quoted literal as code when its content
// starts like code: DECLARE/BEGIN is formatted as a PL/pgSQL block and
// SELECT/INSERT/UPDATE/DELETE/WITH as DML. Any other content (SQL fragments,
// prose, other languages) is left exactly as written, as is the literal if the
// formatted result is not provably equivalent.
func formatEmbedded(lit string, st config.Style) string {
	open, inner, close, ok := splitDollarQuote(lit)
	if !ok {
		return lit
	}
	// format()-style templates (%s, %I, %L, %1$s) are text, not SQL: restyling
	// would split the placeholders, and the token-level safety gate can't see it.
	if formatPlaceholder.MatchString(inner) {
		return lit
	}
	var first string
	for _, t := range lexer.Lex(inner) {
		if t.IsTrivia() || t.Kind == lexer.EOF {
			continue
		}
		if t.Kind == lexer.Ident {
			first = lowerASCII(t.Text)
		}
		break
	}
	var res string
	switch first {
	case "declare", "begin":
		res = formatBodyInner(inner, st)
	case "select", "insert", "update", "delete", "with":
		lead := inner[:len(inner)-len(strings.TrimLeft(inner, " \t\r\n"))]
		trail := inner[len(strings.TrimRight(inner, " \t\r\n")):]
		res = lead + strings.TrimSpace(File(parser.Parse(inner), st)) + trail
	default:
		return lit
	}
	if res == inner || !SemanticallyEqual(inner, res) || CommentsPreserved(inner, res) != nil {
		return lit
	}
	return open + res + close
}

var formatPlaceholder = regexp.MustCompile(`%(\d+\$)?-?\d*[sILlx]`)