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

package mathml

import "strings"

// envSpec describes a table environment: its delimiters, its column
// alignment, whether it renders in script size, whether it takes a column
// specification and whether it takes a number argument first.
type envSpec struct {
	fences   [2]string
	align    string
	scripted bool
	colSpec  bool
	number   bool
}

var environments = map[string]envSpec{
	"matrix": {}, "pmatrix": {fences: [2]string{"(", ")"}},
	"bmatrix":     {fences: [2]string{"[", "]"}},
	"Bmatrix":     {fences: [2]string{"{", "}"}},
	"vmatrix":     {fences: [2]string{"|", "|"}},
	"Vmatrix":     {fences: [2]string{"‖", "‖"}},
	"matrix*":     {},
	"pmatrix*":    {fences: [2]string{"(", ")"}},
	"bmatrix*":    {fences: [2]string{"[", "]"}},
	"Bmatrix*":    {fences: [2]string{"{", "}"}},
	"vmatrix*":    {fences: [2]string{"|", "|"}},
	"Vmatrix*":    {fences: [2]string{"‖", "‖"}},
	"smallmatrix": {scripted: true},
	"cases":       {fences: [2]string{"{", ""}, align: "left"},
	"rcases":      {fences: [2]string{"", "}"}, align: "left"},
	"dcases":      {fences: [2]string{"{", ""}, align: "left"},
	"drcases":     {fences: [2]string{"", "}"}, align: "left"},
	"aligned":     {align: "right left"},
	"align":       {align: "right left"},
	"align*":      {align: "right left"},
	"alignedat":   {align: "right left", number: true},
	"alignat":     {align: "right left", number: true},
	"alignat*":    {align: "right left", number: true},
	"split":       {align: "right left"},
	"gather":      {}, "gather*": {},
	"equation": {}, "equation*": {},
	"array":    {colSpec: true},
	"darray":   {colSpec: true},
	"subarray": {colSpec: true, scripted: true},
}

// environment parses a whole \begin{name}...\end{name} construct. The
// cursor sits just after the \begin token.
func (p *parser) environment() *node {
	begin := p.toks[p.pos-1]
	name, ok := p.envName()
	if !ok {
		return errorNode(begin.text)
	}
	spec, supported := environments[name]
	if !supported {
		return p.degradeEnvironment(begin)
	}

	if spec.colSpec {
		raw, ok := p.rawBraced()
		if !ok {
			return errorNode(begin.text)
		}
		_, supported = parseColSpec(raw)
		if !supported {
			return p.degradeEnvironment(begin)
		}
	} else if spec.number {
		if _, ok := p.rawBraced(); !ok {
			return errorNode(begin.text)
		}
	}
	rows := p.tableRows()
	if !p.atCommand("end") {
		p.pos = len(p.toks) - 1
		return errorNode(string(p.src[begin.start:]))
	}
	endStart := p.toks[p.pos].start
	p.pos++
	endName, ok := p.envName()
	if !ok || endName != name {
		p.pos = len(p.toks) - 1
		return errorNode(string(p.src[begin.start:]))
	}
	_ = endStart
	table := buildTable(rows, spec)
	if spec.scripted {
		inner := elA("mstyle", []attribute{{"scriptlevel", "1"}})
		inner.children = []*node{table}
		table = inner
	}
	if spec.fences == [2]string{"", ""} {
		return table
	}
	row := el("mrow")
	row.children = []*node{p.fenceOf(spec.fences[0]), table, p.fenceOf(spec.fences[1])}
	return row
}

// envName reads the environment name in braces.
func (p *parser) envName() (string, bool) {
	raw, ok := p.rawBraced()
	if !ok || raw == "" {
		return "", false
	}
	return raw, true
}

// degradeEnvironment consumes a whole unsupported environment, up to and
// including its matching \end, and degrades it as its verbatim source.
func (p *parser) degradeEnvironment(begin token) *node {
	depth := 1
	end := len(p.src)
	i := p.pos
	for ; i < len(p.toks); i++ {
		t := p.toks[i]
		if t.kind != tokCommand {
			continue
		}
		switch t.text {
		case `\begin`:
			depth++
		case `\end`:
			depth--
			if depth == 0 {
				end = t.end
				// include the name argument of \end
				if i+2 < len(p.toks) && p.toks[i+1].kind == tokLBrace && p.toks[i+2].kind == tokRBrace {
					end = p.toks[i+2].end
					i += 2
				}
				i++
				p.pos = i
				return errorNode(string(p.src[begin.start:end]))
			}
		}
	}
	p.pos = len(p.toks) - 1
	return errorNode(string(p.src[begin.start:]))
}

// tableRows parses the rows of a table, each row a slice of cells, until
// the \end or the end of input.
func (p *parser) tableRows() [][]*node {
	var rows [][]*node
	for {
		var cells []*node
		for {
			nodes := p.sequence(true)
			cell := el("mtd")
			cell.children = nodes
			cells = append(cells, cell)
			if p.at(tokAmpersand) {
				p.pos++
				continue
			}
			break
		}
		rows = append(rows, cells)
		if isRowEnd(p.cur()) {
			p.pos++
			p.rowSpacing()
			if p.at(tokEOF) || p.atCommand("end") {
				break
			}
			continue
		}
		break
	}
	return rows
}

// rowSpacing skips the optional bracket after a row separator.
func (p *parser) rowSpacing() {
	if p.at(tokChar) && p.cur().text == "[" {
		for {
			t := p.cur()
			p.pos++
			if t.kind == tokEOF || t.kind == tokChar && t.text == "]" {
				return
			}
		}
	}
}

// parseColSpec reads an array column specification: alignment letters and
// vertical rules, nothing else.
func parseColSpec(raw string) (int, bool) {
	count := 0
	for _, c := range raw {
		switch c {
		case 'l', 'c', 'r':
			count++
		case '|', ' ', '\t':
		default:
			return 0, false
		}
	}
	if count == 0 {
		return 0, false
	}
	return count, true
}

// buildTable assembles the mtable with its alignment attributes. The
// "right left" alignment alternates over the widest row.
func buildTable(rows [][]*node, spec envSpec) *node {
	table := el("mtable")
	cols := 0
	for _, row := range rows {
		if len(row) > cols {
			cols = len(row)
		}
	}
	switch {
	case spec.align == "left" && cols > 0:
		table.attrs = append(table.attrs, attribute{"columnalign", "left"})
	case spec.align == "right left" && cols > 0:
		var b strings.Builder
		for i := range cols {
			if i > 0 {
				b.WriteString(" ")
			}
			if i%2 == 0 {
				b.WriteString("right")
			} else {
				b.WriteString("left")
			}
		}
		table.attrs = append(table.attrs, attribute{"columnalign", b.String()})
	}
	for _, row := range rows {
		tr := el("mtr")
		tr.children = row
		table.children = append(table.children, tr)
	}
	return table
}

// substack renders the rows of a \substack argument in script size.
func (p *parser) substack(raw string) *node {
	q := newParser([]byte(raw), false)
	rows := q.tableRows()
	table := buildTable(rows, envSpec{})
	inner := elA("mstyle", []attribute{{"scriptlevel", "1"}})
	inner.children = []*node{table}
	return inner
}