PgTidy/pkg/format/dml.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"
)

// isDMLStart reports whether toks begins with a DML statement keyword.
func isDMLStart(toks []cst.Tok) bool {
	if len(toks) == 0 || toks[0].Tok.Kind != lexer.Ident {
		return false
	}
	switch lowerASCII(toks[0].Tok.Text) {
	case "select", "insert", "update", "delete", "with":
		return true
	}
	return false
}

// dmlSeg is one major clause of a DML statement.
type dmlSeg struct {
	kw   []cst.Tok // clause keyword tokens (possibly multi-word)
	body []cst.Tok // remaining tokens up to the next clause boundary
}

// formatDML formats a top-level DML statement from its token slice, applying
// keyword casing, clause-per-line layout, and leading-comma column lists for
// SELECT and UPDATE SET clauses. Falls back to verbatim on comment-heavy input.
func formatDML(toks []cst.Tok, st config.Style) string {
	// Strip the trailing semicolon so segments don't see it.
	semi := ""
	if n := len(toks); n > 0 && toks[n-1].Tok.Kind == lexer.Semicolon {
		semi = ";"
		toks = toks[:n-1]
	}

	segs := segmentDML(toks)
	if len(segs) == 0 {
		return verbatimSpan(toks) + semi
	}

	nl := st.Newline
	var b strings.Builder
	for i, seg := range segs {
		if i > 0 {
			b.WriteString(nl)
		}
		b.WriteString(dmlSegText(seg, st))
	}
	b.WriteString(semi)
	return b.String()
}

// segmentDML splits toks into clause segments at depth-0 clause boundaries.
// Tokens inside parentheses (depth > 0) are never treated as clause starters,
// so subqueries and function calls are kept intact.
func segmentDML(toks []cst.Tok) []dmlSeg {
	var segs []dmlSeg
	depth := 0
	segStart := 0
	kwEnd := 0
	started := false

	flush := func(end int) {
		if !started || end <= segStart {
			return
		}
		segs = append(segs, dmlSeg{
			kw:   toks[segStart:kwEnd],
			body: toks[kwEnd:end],
		})
	}

	i := 0
	for i < len(toks) {
		t := toks[i]
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
			i++
			continue
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
			i++
			continue
		}

		if depth == 0 && t.Tok.Kind == lexer.Ident {
			kw := lowerASCII(t.Tok.Text)
			if dmlIsClauseKw(kw, toks, i) {
				flush(i)
				started = true
				segStart = i
				i = dmlConsumeKw(toks, i)
				kwEnd = i
				continue
			}
		}
		i++
	}
	flush(len(toks))
	return segs
}

// dmlIsClauseKw reports whether the keyword at toks[i] starts a new DML clause.
func dmlIsClauseKw(kw string, toks []cst.Tok, i int) bool {
	switch kw {
	case "select", "from", "where", "having", "limit", "offset",
		"returning", "with", "into", "values", "set",
		"union", "intersect", "except",
		"insert", "update", "delete",
		"join", "left", "right", "inner", "full", "cross", "natural":
		return true
	case "group", "order":
		return i+1 < len(toks) && toks[i+1].Is("by")
	case "on":
		return i+1 < len(toks) && toks[i+1].Is("conflict")
	}
	return false
}

// dmlConsumeKw advances past multi-word clause keywords (e.g. GROUP BY,
// LEFT OUTER JOIN, INSERT INTO, DELETE FROM) and returns the new index.
func dmlConsumeKw(toks []cst.Tok, i int) int {
	if i >= len(toks) {
		return i
	}
	kw := lowerASCII(toks[i].Tok.Text)
	i++
	switch kw {
	case "group", "order":
		if i < len(toks) && toks[i].Is("by") {
			i++
		}
	case "left", "right", "full":
		if i < len(toks) && toks[i].Is("outer") {
			i++
		}
		if i < len(toks) && toks[i].Is("join") {
			i++
		}
	case "inner", "cross", "natural":
		if i < len(toks) && toks[i].Is("join") {
			i++
		}
	case "on":
		if i < len(toks) && toks[i].Is("conflict") {
			i++
		}
	case "delete":
		// DELETE FROM — consume the FROM so it isn't treated as a separate clause.
		if i < len(toks) && toks[i].Is("from") {
			i++
		}
	case "insert":
		// INSERT INTO — consume INTO.
		if i < len(toks) && toks[i].Is("into") {
			i++
		}
	}
	return i
}

// dmlSegText formats one DML clause segment into a text line (or lines for
// column-list clauses and WITH bodies).
func dmlSegText(seg dmlSeg, st config.Style) string {
	kwText := dmlInline(seg.kw, st)
	kw := ""
	if len(seg.kw) > 0 {
		kw = lowerASCII(seg.kw[0].Tok.Text)
	}

	switch kw {
	case "select", "returning":
		items := dmlSplitCommas(seg.body)
		return dmlColListSelect(kwText, items, st)
	case "set":
		items := dmlSplitCommas(seg.body)
		return dmlColListSet(kwText, items, st)
	case "values":
		return dmlValuesClause(kwText, seg.body, st)
	case "where":
		return dmlWhereClause(kwText, seg.body, st)
	case "join", "left", "right", "inner", "full", "cross", "natural":
		return dmlJoinClause(kwText, seg.body, st)
	case "with":
		return formatWithBody(kwText, seg.body, st)
	default:
		body := dmlInline(seg.body, st)
		if body == "" {
			return kwText
		}
		return kwText + " " + body
	}
}

// dmlJoinClause formats a JOIN clause, applying indent_join when configured.
func dmlJoinClause(kwText string, body []cst.Tok, st config.Style) string {
	text := dmlInline(body, st)
	line := kwText
	if text != "" {
		line += " " + text
	}
	if !st.IndentJoin {
		return line
	}
	indent := strings.Repeat(st.Indent, st.JoinIndentSize)
	nl := st.Newline
	var b strings.Builder
	for i, part := range strings.Split(line, nl) {
		if i > 0 {
			b.WriteString(nl)
		}
		b.WriteString(indent)
		b.WriteString(part)
	}
	return b.String()
}

// dmlWhereClause formats a WHERE clause, splitting AND/OR conditions per
// the where_wrap and where_and_or_indent settings.
func dmlWhereClause(kwText string, body []cst.Tok, st config.Style) string {
	if st.WhereWrap == config.WrapNever {
		text := dmlInline(body, st)
		if text == "" {
			return kwText
		}
		return kwText + " " + text
	}

	// Split at depth-0 AND/OR.
	conditions := dmlSplitAndOr(body)
	if len(conditions) <= 1 {
		text := dmlInline(body, st)
		if text == "" {
			return kwText
		}
		return kwText + " " + text
	}

	nl := st.Newline
	var b strings.Builder
	b.WriteString(kwText)
	for i, cond := range conditions {
		b.WriteString(nl)
		text := dmlInline(cond, st)
		prefix := ""
		if st.WhereAndOrIndent {
			prefix = st.Indent
		}
		if i == 0 {
			prefix += "  " // align with AND/OR token width
		}
		writeListItem(&b, prefix, prefix, text, false, nl)
	}
	return b.String()
}

// dmlSplitAndOr splits toks at depth-0 AND/OR tokens, keeping the AND/OR with
// the following condition.
func dmlSplitAndOr(toks []cst.Tok) [][]cst.Tok {
	var result [][]cst.Tok
	depth := 0
	start := 0
	for i, t := range toks {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 && t.Tok.Kind == lexer.Ident {
			low := lowerASCII(t.Tok.Text)
			if (low == "and" || low == "or") && i > start {
				result = append(result, toks[start:i])
				start = i
			}
		}
	}
	result = append(result, toks[start:])
	return result
}

// formatWithBody formats the body of a WITH clause by splitting CTE definitions
// at depth-0 commas and formatting the subquery inside each AS (...) block.
func formatWithBody(kwText string, body []cst.Tok, st config.Style) string {
	nl := st.Newline

	cteDefs := dmlSplitCommas(body)
	// Filter spurious empty items.
	var kept [][]cst.Tok
	for _, d := range cteDefs {
		if len(d) > 0 {
			kept = append(kept, d)
		}
	}
	cteDefs = kept

	switch len(cteDefs) {
	case 0:
		return kwText
	case 1:
		return kwText + " " + formatCTEDef(cteDefs[0], st)
	default:
		// Multiple CTEs: one per line with the configured comma style.
		first := st.Indent + " "
		cont := st.Indent + ","
		contPad := strings.Repeat(" ", len(cont)) // same width as cont, no comma

		var b strings.Builder
		b.WriteString(kwText)
		for i, cteDef := range cteDefs {
			b.WriteString(nl)

			var headPfx, tailPfx string
			if i == 0 || st.Commas != config.CommaLeading {
				headPfx = first
				tailPfx = first
			} else {
				headPfx = cont
				tailPfx = contPad
			}

			cteText := formatCTEDef(cteDef, st)
			cteLines := strings.Split(cteText, nl)
			for j, line := range cteLines {
				if j > 0 {
					b.WriteString(nl)
					b.WriteString(tailPfx)
				} else {
					b.WriteString(headPfx)
				}
				b.WriteString(line)
			}
		}
		return b.String()
	}
}

// formatCTEDef formats one CTE definition of the form:
//
//	name [column_list] AS [NOT] [MATERIALIZED] (subquery)
//
// The subquery is formatted as DML, indented by st.Indent inside the parentheses.
// Falls back to dmlInline if the expected structure is not found.
func formatCTEDef(toks []cst.Tok, st config.Style) string {
	nl := st.Newline

	// Find the AS keyword at depth 0.
	asIdx := dmlKeywordIdx(toks, 0, "as")
	if asIdx < 0 {
		return dmlInline(toks, st)
	}

	// Find the opening '(' after AS (may be preceded by NOT / MATERIALIZED).
	parenOpen := -1
	for i := asIdx + 1; i < len(toks); i++ {
		if toks[i].Tok.Kind == lexer.LParen {
			parenOpen = i
			break
		}
		if toks[i].Tok.Kind != lexer.Ident {
			// Unexpected token before '(' — fall back.
			break
		}
	}
	if parenOpen < 0 {
		return dmlInline(toks, st)
	}

	// Find the matching ')'.
	parenClose := dmlMatchParen(toks, parenOpen)
	if parenClose < 0 {
		return dmlInline(toks, st)
	}

	// Format the header (name, optional column list, AS, optional MATERIALIZED).
	header := dmlInline(toks[:parenOpen], st)

	// Format and wrap the subquery per subquery_content/subquery_closing.
	// The "AS (" space is standard CTE syntax and independent of
	// subquery_space_before_paren; only subquery_opening's newline choice
	// applies here.
	subToks := toks[parenOpen+1 : parenClose]
	sep := " "
	if st.SubqueryOpening == config.PlacementNewLine {
		sep = nl
	}
	return header + sep + dmlWrapSubquery(subToks, st)
}

// dmlKeywordIdx returns the index of the first token equal to kw at paren depth 0,
// starting from `from`. Returns -1 if not found.
func dmlKeywordIdx(toks []cst.Tok, from int, kw string) int {
	depth := 0
	for i := from; i < len(toks); i++ {
		switch toks[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 && toks[i].Is(kw) {
			return i
		}
	}
	return -1
}

// dmlMatchParen returns the index of the ')' matching the '(' at toks[open].
// Returns -1 if no matching paren is found.
func dmlMatchParen(toks []cst.Tok, open int) int {
	depth := 1
	for i := open + 1; i < len(toks); i++ {
		switch toks[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen:
			depth--
			if depth == 0 {
				return i
			}
		case lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
	}
	return -1
}

// dmlIsSubqueryOpen reports whether toks[i] is a '(' immediately followed by
// SELECT or WITH — i.e. it opens a subquery (derived table, scalar subquery,
// or an IN/EXISTS/ANY/ALL/ARRAY(...) subquery), as opposed to a function-call
// argument list, a value tuple, or a grouping paren.
func dmlIsSubqueryOpen(toks []cst.Tok, i int) bool {
	if toks[i].Tok.Kind != lexer.LParen {
		return false
	}
	j := i + 1
	if j >= len(toks) || toks[j].Tok.Kind != lexer.Ident {
		return false
	}
	switch lowerASCII(toks[j].Tok.Text) {
	case "select", "with":
		return true
	}
	return false
}

// dmlWrapSubquery formats a subquery's inner tokens (excluding the enclosing
// parens) as DML and wraps them in "(" … ")" per the subquery_content and
// subquery_closing settings. The result is rendered relative to column 0;
// callers that splice it mid-line are responsible for re-indenting any
// continuation lines to the surrounding context.
func dmlWrapSubquery(inner []cst.Tok, st config.Style) string {
	nl := st.Newline
	sub := strings.TrimRight(formatDML(inner, st), nl)
	if sub == "" {
		return "()"
	}
	lines := strings.Split(sub, nl)

	var b strings.Builder
	b.WriteString("(")
	for i, line := range lines {
		if i == 0 && st.SubqueryContent != config.PlacementNewLine {
			b.WriteString(line)
			continue
		}
		b.WriteString(nl)
		if line != "" {
			b.WriteString(st.Indent)
		}
		b.WriteString(line)
	}
	if st.SubqueryClosing == config.PlacementNewLine {
		b.WriteString(nl)
	}
	b.WriteString(")")
	return b.String()
}

// dmlInline renders toks on one line with keyword casing and proper spacing.
// If toks[1:] contains comment trivia the function falls back to verbatimSpan
// so no comment is lost. Subquery parens and CASE…END expressions embedded
// anywhere in toks are recursively formatted and spliced in.
func dmlInline(toks []cst.Tok, st config.Style) string {
	if len(toks) == 0 {
		return ""
	}
	if anyComment(toks[1:]) {
		return verbatimSpan(toks)
	}
	nl := st.Newline
	var b strings.Builder
	i := 0
	for i < len(toks) {
		t := toks[i]

		if dmlIsSubqueryOpen(toks, i) {
			if closeIdx := dmlMatchParen(toks, i); closeIdx > i {
				dmlWriteSubquerySep(&b, toks, i, st, nl)
				b.WriteString(dmlWrapSubquery(toks[i+1:closeIdx], st))
				i = closeIdx + 1
				continue
			}
		}

		if dmlIsCaseStart(t) {
			if endIdx := dmlMatchCaseEnd(toks, i); endIdx > i {
				if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
					b.WriteByte(' ')
				}
				b.WriteString(dmlFormatCase(toks[i:endIdx+1], st))
				i = endIdx + 1
				continue
			}
		}

		if i > 0 {
			space := needSpace(toks[i-1].Tok, t.Tok) && !isPctTypeBoundary(toks, i)
			if !space && st.RecordSpaceBeforeParen && t.Tok.Kind == lexer.LParen &&
				toks[i-1].Tok.Kind == lexer.Ident && lowerASCII(toks[i-1].Tok.Text) == "row" {
				space = true
			}
			if space {
				b.WriteByte(' ')
			}
		}
		// Space after comma in calls: func(a, b) vs func(a,b).
		if st.SpaceAfterCommaInCalls && i > 0 && toks[i-1].Tok.Kind == lexer.Comma {
			// Only inside parens (caller manages this at depth > 0, but we add space
			// when the comma is not a clause-level comma — heuristic: always add).
			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, st))
		i++
	}
	return b.String()
}

// dmlWriteSubquerySep writes the separator between the token preceding a
// subquery-opening '(' at toks[i] and the '(' itself, honoring
// subquery_opening (same_line|new_line) and subquery_space_before_paren.
func dmlWriteSubquerySep(b *strings.Builder, toks []cst.Tok, i int, st config.Style, nl string) {
	if i == 0 {
		return
	}
	if st.SubqueryOpening == config.PlacementNewLine {
		b.WriteString(nl)
		return
	}
	space := needSpace(toks[i-1].Tok, toks[i].Tok) && !isPctTypeBoundary(toks, i)
	if !space && st.SubquerySpaceBeforeParen {
		space = true
	}
	if space {
		b.WriteByte(' ')
	}
}

// dmlSplitCommas splits toks at depth-0 commas and returns the items between
// them (the comma tokens themselves are discarded).
func dmlSplitCommas(toks []cst.Tok) [][]cst.Tok {
	var items [][]cst.Tok
	depth := 0
	start := 0
	for i, t := range toks {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		case lexer.Comma:
			if depth == 0 {
				items = append(items, toks[start:i])
				start = i + 1
			}
		}
	}
	// Remaining tokens after the last comma (or all tokens if no comma found).
	items = append(items, toks[start:])
	return items
}

// filterEmpty drops empty token slices (spurious items from a trailing
// comma or similar).
func filterEmpty(items [][]cst.Tok) [][]cst.Tok {
	var kept [][]cst.Tok
	for _, item := range items {
		if len(item) > 0 {
			kept = append(kept, item)
		}
	}
	return kept
}

// dmlCommaList renders texts as a one-item-per-line list under kwText, using
// leading or trailing commas per st.Commas. Items whose rendered text spans
// multiple lines (e.g. an embedded subquery or wrapped CASE) have their
// continuation lines re-indented to align under the item's first line.
func dmlCommaList(kwText string, texts []string, st config.Style) string {
	switch len(texts) {
	case 0:
		return kwText
	case 1:
		if texts[0] == "" {
			return kwText
		}
		return kwText + " " + texts[0]
	}

	nl := st.Newline
	first := st.Indent + " "
	cont := st.Indent + ","
	contPad := strings.Repeat(" ", len(cont))

	var b strings.Builder
	b.WriteString(kwText)
	for i, text := range texts {
		b.WriteString(nl)
		trailingComma := st.Commas == config.CommaTrailing && i < len(texts)-1
		if i == 0 || st.Commas != config.CommaLeading {
			writeListItem(&b, first, first, text, trailingComma, nl)
		} else {
			writeListItem(&b, cont, contPad, text, false, nl)
		}
	}
	return b.String()
}

// writeListItem writes text prefixed with headPfx (its first line) and
// tailPfx (any continuation lines), optionally followed by a trailing comma.
func writeListItem(b *strings.Builder, headPfx, tailPfx, text string, trailingComma bool, nl string) {
	for j, line := range strings.Split(text, nl) {
		if j > 0 {
			b.WriteString(nl)
			if line != "" {
				b.WriteString(tailPfx)
			}
		} else {
			b.WriteString(headPfx)
		}
		b.WriteString(line)
	}
	if trailingComma {
		b.WriteString(",")
	}
}

// dmlColListSelect formats a SELECT / RETURNING column list with optional
// align_columns and select_align_as settings.
func dmlColListSelect(kwText string, items [][]cst.Tok, st config.Style) string {
	items = filterEmpty(items)
	if len(items) == 1 {
		body := dmlInline(items[0], st)
		if body == "" {
			return kwText
		}
		return kwText + " " + body
	}

	texts := make([]string, len(items))
	for i, item := range items {
		texts[i] = dmlInline(item, st)
	}

	// align_columns / select_align_as: pad expressions so AS and aliases align.
	if (st.AlignColumns || st.SelectAlignAs) && len(texts) > 1 {
		texts = alignSelectItems(texts, st)
	}

	return dmlCommaList(kwText, texts, st)
}

// dmlColListSet formats an UPDATE SET column list with optional set_align_equal.
func dmlColListSet(kwText string, items [][]cst.Tok, st config.Style) string {
	items = filterEmpty(items)
	if len(items) == 1 {
		body := dmlInline(items[0], st)
		if body == "" {
			return kwText
		}
		return kwText + " " + body
	}

	texts := make([]string, len(items))
	for i, item := range items {
		texts[i] = dmlInline(item, st)
	}

	// set_align_equal: pad lhs so = signs align.
	if st.SetAlignEqual && len(texts) > 1 {
		texts = alignSetItems(texts)
	}

	return dmlCommaList(kwText, texts, st)
}

// dmlValuesClause formats a VALUES clause. When insert_collapse_values is
// true (the default), multiple rows stay packed onto one line, matching the
// pre-existing flat rendering. When false, each row gets its own line.
func dmlValuesClause(kwText string, body []cst.Tok, st config.Style) string {
	rows := filterEmpty(dmlSplitCommas(body))

	if len(rows) <= 1 || st.InsertCollapseValues {
		text := dmlInline(body, st)
		if text == "" {
			return kwText
		}
		return kwText + " " + text
	}

	texts := make([]string, len(rows))
	for i, row := range rows {
		texts[i] = dmlInline(row, st)
	}
	return dmlCommaList(kwText, texts, st)
}

// alignSelectItems pads SELECT list item expressions so that AS keywords and
// alias names align vertically.
func alignSelectItems(texts []string, st config.Style) []string {
	// Split each text into (expr, " AS ", alias) or keep as-is.
	type part struct {
		expr, alias string
		hasAs       bool
	}
	parts := make([]part, len(texts))
	maxExpr := 0
	for i, t := range texts {
		// Find " AS " or " as " (case-insensitive).
		if idx := findAsIndex(t); idx >= 0 {
			parts[i] = part{expr: t[:idx], alias: t[idx:], hasAs: true}
			if l := len(t[:idx]); l > maxExpr {
				maxExpr = l
			}
		} else {
			parts[i] = part{expr: t}
			if st.AlignColumns {
				if l := len(t); l > maxExpr {
					maxExpr = l
				}
			}
		}
	}

	out := make([]string, len(texts))
	for i, p := range parts {
		if !p.hasAs || maxExpr == 0 {
			out[i] = texts[i]
			continue
		}
		pad := strings.Repeat(" ", maxExpr-len(p.expr))
		out[i] = p.expr + pad + p.alias
	}
	return out
}

// findAsIndex returns the byte index of " AS " (case-insensitive) in s,
// or -1 if not present at depth 0.
func findAsIndex(s string) int {
	low := lowerASCII(s)
	// Look for " as " boundary.
	for i := 0; i < len(low)-3; i++ {
		if low[i] == ' ' && low[i+1] == 'a' && low[i+2] == 's' && low[i+3] == ' ' {
			return i + 1 // index of 'a'
		}
	}
	return -1
}

// alignSetItems pads SET assignment lhs values so that = signs align.
func alignSetItems(texts []string) []string {
	maxLhs := 0
	lhsWidths := make([]int, len(texts))
	for i, t := range texts {
		idx := strings.Index(t, " = ")
		if idx < 0 {
			idx = strings.Index(t, "=")
		}
		if idx >= 0 {
			lhsWidths[i] = idx
			if idx > maxLhs {
				maxLhs = idx
			}
		}
	}
	if maxLhs == 0 {
		return texts
	}
	out := make([]string, len(texts))
	for i, t := range texts {
		if lhsWidths[i] == 0 || lhsWidths[i] == maxLhs {
			out[i] = t
			continue
		}
		idx := lhsWidths[i]
		pad := strings.Repeat(" ", maxLhs-idx)
		out[i] = t[:idx] + pad + t[idx:]
	}
	return out
}

// caseCollapseWidth is the inline-length threshold under which case_collapse
// keeps a CASE expression on one line even when case_when_wrap is set.
const caseCollapseWidth = 60

// dmlIsCaseStart reports whether t is a CASE keyword token.
func dmlIsCaseStart(t cst.Tok) bool {
	return t.Tok.Kind == lexer.Ident && lowerASCII(t.Tok.Text) == "case"
}

// dmlMatchCaseEnd returns the index of the END token that closes the CASE
// token at toks[start], accounting for nested CASE…END and paren depth.
// Returns -1 if no matching END is found.
func dmlMatchCaseEnd(toks []cst.Tok, start int) int {
	depth := 0
	caseDepth := 1
	for i := start + 1; i < len(toks); i++ {
		switch toks[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
			continue
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
			continue
		}
		if depth != 0 || toks[i].Tok.Kind != lexer.Ident {
			continue
		}
		switch lowerASCII(toks[i].Tok.Text) {
		case "case":
			caseDepth++
		case "end":
			caseDepth--
			if caseDepth == 0 {
				return i
			}
		}
	}
	return -1
}

// caseSeg is one part of a CASE expression's body: the optional leading
// operand (kw == nil), or a WHEN/THEN/ELSE-led span.
type caseSeg struct {
	kw   *cst.Tok
	toks []cst.Tok
}

// dmlSplitCase splits a CASE expression's body (the tokens strictly between
// CASE and its matching END) into operand/when/then/else segments at
// depth-0 boundaries, skipping over any nested CASE…END.
func dmlSplitCase(body []cst.Tok) []caseSeg {
	var segs []caseSeg
	depth := 0
	caseDepth := 0
	start := 0
	var curKw *cst.Tok
	flush := func(end int) {
		if end > start {
			segs = append(segs, caseSeg{kw: curKw, toks: body[start:end]})
		}
	}
	for i := range body {
		t := body[i]
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
			continue
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
			continue
		}
		if depth != 0 || t.Tok.Kind != lexer.Ident {
			continue
		}
		switch lowerASCII(t.Tok.Text) {
		case "case":
			caseDepth++
		case "end":
			if caseDepth > 0 {
				caseDepth--
			}
		case "when", "then", "else":
			if caseDepth == 0 {
				flush(i)
				start = i + 1
				kw := body[i]
				curKw = &kw
			}
		}
	}
	flush(len(body))
	return segs
}

// whenThen is one rendered WHEN … THEN … branch of a CASE expression.
type whenThen struct {
	whenKw, cond, thenKw, then string
}

// dmlFormatCase renders a CASE…END expression honoring case_when_wrap,
// case_end, and case_collapse. toks[0] must be CASE and toks[len(toks)-1]
// its matching END.
func dmlFormatCase(toks []cst.Tok, st config.Style) string {
	body := toks[1 : len(toks)-1]
	segs := dmlSplitCase(body)

	operand := ""
	var whens []whenThen
	elseKw, elseText := "", ""
	haveElse := false
	pendingWhenKw, pendingCond := "", ""

	for _, s := range segs {
		text := dmlInline(s.toks, st)
		if s.kw == nil {
			operand = text
			continue
		}
		kwText := caseText(s.kw.Tok, st)
		switch lowerASCII(s.kw.Tok.Text) {
		case "when":
			pendingWhenKw, pendingCond = kwText, text
		case "then":
			whens = append(whens, whenThen{whenKw: pendingWhenKw, cond: pendingCond, thenKw: kwText, then: text})
		case "else":
			elseKw, elseText, haveElse = kwText, text, true
		}
	}

	caseKw := caseText(toks[0].Tok, st)
	endKw := caseText(toks[len(toks)-1].Tok, st)

	inline := dmlCaseInline(caseKw, operand, whens, elseKw, elseText, haveElse, endKw)
	if !st.CaseWhenWrap {
		return inline
	}
	if st.CaseCollapse && len(inline) <= caseCollapseWidth {
		return inline
	}
	return dmlCaseWrapped(caseKw, operand, whens, elseKw, elseText, haveElse, endKw, st)
}

// dmlCaseInline renders a CASE expression on a single line.
func dmlCaseInline(caseKw, operand string, whens []whenThen, elseKw, elseText string, haveElse bool, endKw string) string {
	var b strings.Builder
	b.WriteString(caseKw)
	if operand != "" {
		b.WriteByte(' ')
		b.WriteString(operand)
	}
	for _, w := range whens {
		b.WriteByte(' ')
		b.WriteString(w.whenKw)
		if w.cond != "" {
			b.WriteByte(' ')
			b.WriteString(w.cond)
		}
		b.WriteByte(' ')
		b.WriteString(w.thenKw)
		if w.then != "" {
			b.WriteByte(' ')
			b.WriteString(w.then)
		}
	}
	if haveElse {
		b.WriteByte(' ')
		b.WriteString(elseKw)
		if elseText != "" {
			b.WriteByte(' ')
			b.WriteString(elseText)
		}
	}
	b.WriteByte(' ')
	b.WriteString(endKw)
	return b.String()
}

// dmlCaseWrapped renders a CASE expression with each WHEN … THEN branch (and
// ELSE) on its own line, per case_end for the closing END's placement.
func dmlCaseWrapped(caseKw, operand string, whens []whenThen, elseKw, elseText string, haveElse bool, endKw string, st config.Style) string {
	nl := st.Newline
	indent := st.Indent

	var b strings.Builder
	b.WriteString(caseKw)
	if operand != "" {
		b.WriteByte(' ')
		b.WriteString(operand)
	}
	for _, w := range whens {
		b.WriteString(nl)
		b.WriteString(indent)
		b.WriteString(w.whenKw)
		if w.cond != "" {
			b.WriteByte(' ')
			b.WriteString(w.cond)
		}
		b.WriteByte(' ')
		b.WriteString(w.thenKw)
		if w.then != "" {
			b.WriteByte(' ')
			b.WriteString(w.then)
		}
	}
	if haveElse {
		b.WriteString(nl)
		b.WriteString(indent)
		b.WriteString(elseKw)
		if elseText != "" {
			b.WriteByte(' ')
			b.WriteString(elseText)
		}
	}
	if st.CaseEnd == config.PlacementNewLine {
		b.WriteString(nl)
		b.WriteString(endKw)
	} else {
		b.WriteByte(' ')
		b.WriteString(endKw)
	}
	return b.String()
}