// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0

package markdown

import (
	"regexp"
	"strconv"
	"strings"

	"sourcedock.dev/petrbalvin/scriptorium"
)

// The mathematics pass turns TeX runs into MathML: $$…$$ occupying a
// whole line becomes a display block, $…$ within one line becomes an
// inline element. scriptorium's Markdown grammar knows nothing about
// dollars, so the runs are lifted out of the source before rendering,
// each replaced by a placeholder of two private-use runes around its
// index, and the MathML is spliced back into the rendered HTML in the
// placeholder's place. A run outside the mappable surface is not
// refused: scriptorium degrades it in place, its verbatim source inside
// an merror element, so nothing is silently mistranslated.

// mathSpan is one equation lifted out of the source.
type mathSpan struct {
	display bool
	src     string
}

// The placeholder runes come from Unicode's private-use area, so no
// authored body contains them; a body that somehow does is left without
// mathematics rather than spliced into the wrong place.
const (
	sentinelOpen  = '\uE000'
	sentinelClose = '\uE001'
)

func mathToken(i int) string {
	return string(sentinelOpen) + strconv.Itoa(i) + string(sentinelClose)
}

// lineKind tells how the previous emitted line ended, which is what the
// display rules need: an equation that follows text needs a separating
// blank line, or the renderer keeps it inside the paragraph above it.
type lineKind int

const (
	prevBlank lineKind = iota
	prevText
	prevDisplay
)

// extractMath rewrites the source with placeholders and returns the
// equations in the order their placeholders appear.
func extractMath(src string) (string, []mathSpan) {
	if strings.ContainsRune(src, sentinelOpen) || strings.ContainsRune(src, sentinelClose) {
		return src, nil
	}
	var spans []mathSpan
	next := func(display bool, src string) string {
		spans = append(spans, mathSpan{display: display, src: src})
		return mathToken(len(spans) - 1)
	}
	var out strings.Builder
	emit := func(line string) {
		out.WriteString(line)
		out.WriteByte('\n')
	}

	fenceChar := byte(0)
	fenceLen := 0
	inCode := false
	inHTML := false
	pendingTick := 0
	prev := prevBlank

	// The lines are indexed rather than streamed: the multi-line display
	// collector looks ahead from the line it stands on.
	lines := strings.Split(src, "\n")
	for i := 0; i < len(lines); i++ {
		line := lines[i]
		indent := len(line) - len(strings.TrimLeft(line, " \t"))
		blank := strings.TrimSpace(line) == ""

		// A fenced code block passes every line through verbatim.
		if fenceChar != 0 {
			emit(line)
			if isClosingFence(line, fenceChar, fenceLen) {
				fenceChar = 0
				prev = prevBlank
			}
			continue
		}

		// The code-span cover of the line, carrying any run an earlier
		// line left open. A span that has not closed keeps the whole
		// line out of every other rule.
		covered, open := codeSpans(line, pendingTick)
		if open > 0 {
			emit(line)
			pendingTick = open
			prev = prevText
			continue
		}
		pendingTick = 0
		startsCovered := len(covered) > 0 && covered[0]

		if !startsCovered {
			if c, n, ok := openingFence(line); ok {
				fenceChar, fenceLen = c, n
				emit(line)
				continue
			}
		}

		// An indented code block: entered from a blank line, left by the
		// first line that is blank or carries less indentation.
		if inCode {
			if blank || indent >= 4 {
				emit(line)
				continue
			}
			inCode = false
		} else if indent >= 4 && prev == prevBlank && !blank {
			inCode = true
			emit(line)
			continue
		}

		// A raw HTML block runs to the next blank line, and its dollars
		// are markup, not mathematics.
		if inHTML {
			emit(line)
			if blank {
				inHTML = false
				prev = prevBlank
			}
			continue
		}
		if !startsCovered && startsHTMLBlock(line) {
			inHTML = true
			emit(line)
			continue
		}

		if blank {
			emit(line)
			prev = prevBlank
			continue
		}

		quotePrefix, rest := stripQuoteMarkers(line)
		listPrefix, item, inList := stripListMarker(rest)

		if inner, ok := displayMath(item); ok {
			if !inList && prev != prevBlank {
				emit(strings.TrimRight(quotePrefix, " \t"))
			}
			emit(quotePrefix + listPrefix + next(true, inner))
			prev = prevDisplay
			continue
		}
		if parts, end, ok := collectDisplayBlock(lines, i, item); ok {
			if !inList && prev != prevBlank {
				emit("")
			}
			emit(quotePrefix + listPrefix + next(true, strings.Join(parts, "\n")))
			prev = prevDisplay
			i = end
			continue
		}
		if prev == prevDisplay && !inList {
			emit(strings.TrimRight(quotePrefix, " \t"))
		}
		off := len(quotePrefix) + len(listPrefix)
		emit(quotePrefix + listPrefix + renderInlineMath(item, covered, off, next))
		prev = prevText
	}
	return out.String(), spans
}

// maxMathBlockLines bounds how far a multi-line display block may reach
// for its closing $$. A block the author never closed then falls back to
// literal text instead of swallowing the rest of the body.
const maxMathBlockLines = 64

// collectDisplayBlock looks ahead from lines[i], whose content opens a
// $$ block it does not close on the same line, for the line that closes
// it. The content lines between become the parts of one display
// equation. The collection stays inside one paragraph: a blank line, a
// code fence, a blockquote marker or the line bound ends the search and
// the block is refused, so a stray $$ stays the text it looks like.
func collectDisplayBlock(lines []string, i int, item string) (parts []string, end int, ok bool) {
	if !strings.HasPrefix(item, "$$") {
		return nil, 0, false
	}
	first := item[2:]
	if strings.Contains(first, "$$") {
		return nil, 0, false
	}
	if strings.TrimSpace(first) != "" {
		parts = append(parts, first)
	}
	for j := i + 1; j < len(lines) && j-i <= maxMathBlockLines; j++ {
		t := strings.TrimSpace(lines[j])
		if t == "" || strings.HasPrefix(t, ">") {
			return nil, 0, false
		}
		if _, _, fenced := openingFence(lines[j]); fenced {
			return nil, 0, false
		}
		if strings.HasSuffix(t, "$$") && backslashRun(t, len(t)-2)%2 == 0 {
			tail := t[:len(t)-2]
			if strings.Contains(tail, "$$") {
				return nil, 0, false
			}
			if strings.TrimSpace(tail) != "" {
				parts = append(parts, tail)
			}
			if len(parts) == 0 || strings.TrimSpace(strings.Join(parts, "")) == "" {
				return nil, 0, false
			}
			return parts, j, true
		}
		parts = append(parts, t)
	}
	return nil, 0, false
}

// renderInlineMath replaces the $…$ runs of one line with placeholders,
// leaving the bytes inside backtick code spans and the dollars written
// \$ alone. The cover was computed for the whole source line, so off
// tells where the line's remaining content begins in it.
func renderInlineMath(line string, covered []bool, off int, next func(bool, string) string) string {
	var b strings.Builder
	i := 0
	for i < len(line) {
		if i+off < len(covered) && covered[i+off] {
			b.WriteByte(line[i])
			i++
			continue
		}
		if line[i] == '$' && backslashRun(line, i)%2 == 0 {
			if value, consumed, ok := matchInlineMath(line[i:]); ok {
				b.WriteString(next(false, value))
				i += consumed
				continue
			}
		}
		b.WriteByte(line[i])
		i++
	}
	return b.String()
}

// displayMath reports whether the whole of a line's content is one
// $$…$$ run, and returns the mathematics between the fences. A run that
// is empty, that hides another $$ or that shares its line with anything
// else is not a display equation.
func displayMath(t string) (string, bool) {
	if len(t) < 5 || !strings.HasPrefix(t, "$$") || !strings.HasSuffix(t, "$$") {
		return "", false
	}
	inner := t[2 : len(t)-2]
	if strings.Contains(inner, "$$") || strings.TrimSpace(inner) == "" {
		return "", false
	}
	return inner, true
}

// matchInlineMath matches one $…$ run at the head of line. The guards
// mirror the shape of real prose: the run must not be empty, may not
// start or end with a space, may not close before a digit, and may not
// contain a dollar inside, so a sentence with two dollar amounts does
// not become a phantom equation, and a run whose opening dollar is a
// closer of an equation broken across lines never swallows the prose
// around it into mathematics.
func matchInlineMath(line string) (value string, consumed int, ok bool) {
	if len(line) < 3 || line[0] != '$' || line[1] == '$' || line[1] == ' ' {
		return "", 0, false
	}
	for i := 1; i < len(line); i++ {
		if line[i] == '\\' {
			i++ // an escaped character is never the closer
			continue
		}
		if line[i] != '$' || i == 1 {
			continue
		}
		if line[i-1] == ' ' {
			continue
		}
		if i+1 < len(line) && isDigit(line[i+1]) {
			continue // currency: the next run of digits belongs outside
		}
		value := line[1:i]
		if strings.ContainsRune(value, '$') {
			return "", 0, false
		}
		return value, i + 1, true
	}
	return "", 0, false
}

func isDigit(b byte) bool { return '0' <= b && b <= '9' }

// backslashRun counts the backslashes immediately before line[i]; an
// odd count means the byte is escaped.
func backslashRun(line string, i int) int {
	n := 0
	for j := i - 1; j >= 0 && line[j] == '\\'; j-- {
		n++
	}
	return n
}

// codeSpans marks the bytes of line that sit inside a backtick code
// span and reports the length of a run the line leaves open. A span
// opens with a run of n backticks and closes with the next run of
// exactly n; a run left open is carried to the next line by the
// pending state, and while it is open no dollar on the line starts
// mathematics.
func codeSpans(line string, pending int) (covered []bool, open int) {
	covered = make([]bool, len(line))
	open = pending
	i := 0
	if pending > 0 {
		end := findTickRun(line, pending)
		if end < 0 {
			for j := range covered {
				covered[j] = true
			}
			return covered, pending
		}
		for j := 0; j < end+pending; j++ {
			covered[j] = true
		}
		i = end + pending
		open = 0
	}
	for i < len(line) {
		if line[i] != '`' {
			i++
			continue
		}
		n := 0
		for i+n < len(line) && line[i+n] == '`' {
			n++
		}
		end := findTickRun(line[i+n:], n)
		if end < 0 {
			if open == 0 {
				open = n
			}
			i += n
			continue
		}
		closeAt := i + n + end
		for j := i; j < closeAt+n; j++ {
			covered[j] = true
		}
		i = closeAt + n
	}
	return covered, open
}

// findTickRun returns the index in line where a run of exactly n
// backticks begins, or -1 when there is none.
func findTickRun(line string, n int) int {
	for i := 0; i < len(line); {
		if line[i] != '`' {
			i++
			continue
		}
		run := 0
		for i+run < len(line) && line[i+run] == '`' {
			run++
		}
		if run == n {
			return i
		}
		i += run
	}
	return -1
}

// openingFence reports whether the line opens a fenced code block, and
// with which character and length.
func openingFence(line string) (byte, int, bool) {
	t := strings.TrimLeft(line, " \t")
	if len(line)-len(t) > 3 {
		return 0, 0, false
	}
	if n := fenceRun(t, '`'); n > 0 {
		return '`', n, true
	}
	if n := fenceRun(t, '~'); n > 0 {
		return '~', n, true
	}
	return 0, 0, false
}

// fenceRun returns the length of a fence of c at the head of t, which
// the rest of the line may follow only with spaces, or 0 when this is
// not a fence.
func fenceRun(t string, c byte) int {
	n := 0
	for n < len(t) && t[n] == c {
		n++
	}
	if n < 3 {
		return 0
	}
	for _, r := range t[n:] {
		if r != ' ' && r != '\t' {
			return 0
		}
	}
	return n
}

// isClosingFence reports whether the line closes an open fence.
func isClosingFence(line string, c byte, n int) bool {
	t := strings.TrimLeft(line, " \t")
	if len(line)-len(t) > 3 {
		return false
	}
	run := 0
	for run < len(t) && t[run] == c {
		run++
	}
	if run < n {
		return false
	}
	for _, r := range t[run:] {
		if r != ' ' && r != '\t' {
			return false
		}
	}
	return true
}

// startsHTMLBlock approximates the CommonMark HTML block: a line that
// opens with a tag, a closing tag, a comment or a declaration starts
// one, and the block then runs to the next blank line.
func startsHTMLBlock(line string) bool {
	t := strings.TrimLeft(line, " \t")
	if len(t) < 2 || t[0] != '<' {
		return false
	}
	c := t[1]
	return c == '/' || c == '!' || c == '?' ||
		('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z')
}

// stripQuoteMarkers removes the blockquote markers from the head of the
// line and returns everything consumed with what remains.
func stripQuoteMarkers(line string) (prefix, rest string) {
	rest = line
	for {
		j := 0
		for j < len(rest) && (rest[j] == ' ' || rest[j] == '\t') {
			j++
		}
		if j >= len(rest) || rest[j] != '>' {
			break
		}
		rest = rest[j+1:]
	}
	return line[:len(line)-len(rest)], rest
}

// stripListMarker removes one list marker from the head of the line, so
// an equation that is a list item's whole content is still recognised
// as display mathematics inside the item.
func stripListMarker(line string) (prefix, rest string, ok bool) {
	j := 0
	for j < len(line) && (line[j] == ' ' || line[j] == '\t') {
		j++
	}
	rest = line[j:]
	if strings.HasPrefix(rest, "- ") || strings.HasPrefix(rest, "* ") || strings.HasPrefix(rest, "+ ") {
		rest = rest[2:]
	} else {
		digits := 0
		for digits < len(rest) && digits < 9 && isDigit(rest[digits]) {
			digits++
		}
		if digits == 0 || digits+1 >= len(rest) {
			return "", line, false
		}
		if (rest[digits] == '.' || rest[digits] == ')') && rest[digits+1] == ' ' {
			rest = rest[digits+2:]
		} else {
			return "", line, false
		}
	}
	rest = strings.TrimLeft(rest, " \t")
	if rest == "" {
		return "", line, false
	}
	return line[:len(line)-len(rest)], rest, true
}

// spliceMath puts the rendered MathML into the HTML in each
// placeholder's place. A placeholder alone in its paragraph becomes the
// display wrapper; any other position takes the math element as it
// stands, which keeps the HTML valid where a display equation shares
// its paragraph with text or sits in a tight list item.
func spliceMath(html string, spans []mathSpan) string {
	for i, span := range spans {
		math := renderMathSpan(span)
		alone := regexp.MustCompile(`(?s)<p>\s*` + mathToken(i) + `\s*</p>`)
		if alone.MatchString(html) {
			html = alone.ReplaceAllString(html, regexpEscapeRepl(math))
			continue
		}
		html = strings.Replace(html, mathToken(i), math, 1)
	}
	return html
}

// regexpEscapeRepl guards the replacement text of ReplaceAllString,
// where a dollar sign would otherwise read as a capture group.
func regexpEscapeRepl(s string) string {
	return strings.ReplaceAll(s, "$", "$$")
}

// renderMathSpan renders one equation through scriptorium. Rendering
// never fails: a construct outside the mappable surface degrades to its
// verbatim source inside an merror element, in place.
func renderMathSpan(span mathSpan) string {
	if span.display {
		return `<div class="math math-display">` + "\n" +
			string(scriptorium.RenderMathDisplay([]byte(span.src))) + "\n</div>"
	}
	return string(scriptorium.RenderMath([]byte(span.src)))
}