package widows

import (
	"bytes"

	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/parser"
	"github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"
)

// nbsp is the space that will not break. Written as bytes because it replaces bytes in the text.
var nbsp = []byte(" ")

// New returns the Markdown extension.
func New() goldmark.Extender { return extension{} }

type extension struct{}

func (extension) Extend(md goldmark.Markdown) {
	md.Parser().AddOptions(parser.WithASTTransformers(util.Prioritized(transformer{}, 200)))
}

// transformer joins the last two words of every paragraph and heading.
//
// It works on the parsed tree rather than on rendered HTML, which is the whole reason this waited for the
// shortcode machinery: a pass over HTML cannot tell prose from an escaped code span, and would happily edit
// the inside of one (content-model.md).
type transformer struct{}

func (transformer) Transform(doc *ast.Document, reader text.Reader, pc parser.Context) {
	source := reader.Source()
	err := ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		switch n.Kind() {
		case ast.KindParagraph, ast.KindHeading:
			join(n, source)
		}
		return ast.WalkContinue, nil
	})
	if err != nil {
		// Walk only fails if this function does, and it does not.
		return
	}
}

// join replaces the last space of a block's final line with a non-breaking one.
//
// Only plain text qualifies. A block ending in a code span, a link or emphasis is left alone: the last
// "word" is then a construct rather than a word, and joining it to the one before would mean reaching inside
// markup for a typographic nicety.
func join(block ast.Node, source []byte) {
	run := finalLine(block)
	if len(run) == 0 {
		return
	}
	var line []byte
	for _, t := range run {
		line = append(line, t.Segment.Value(source)...)
	}
	i := bytes.LastIndexByte(line, ' ')
	if i < 0 || i == len(line)-1 {
		// One word, or a trailing space with nothing after it: nothing to keep together.
		return
	}
	// A String node carries its own bytes, so the text can differ from the source. Editing the segment is
	// impossible — a segment is an offset into bytes every other node shares.
	joined := make([]byte, 0, len(line)+len(nbsp))
	joined = append(joined, line[:i]...)
	joined = append(joined, nbsp...)
	joined = append(joined, line[i+1:]...)
	block.ReplaceChild(block, run[0], ast.NewString(joined))
	for _, t := range run[1:] {
		block.RemoveChild(block, t)
	}
}

// finalLine is the block's last line as the consecutive text nodes that make it up.
//
// Several nodes, not one: another extension may have split the run. The typographer splits it wherever it
// looks for a substitution, so a paragraph ending "hand." arrives as two text nodes and the last of them
// holds no space at all — which is exactly how the first version of this passed its own tests and did
// nothing in the assembled engine.
func finalLine(block ast.Node) []*ast.Text {
	var run []*ast.Text
	last := block.LastChild()
	for n := last; n != nil; n = n.PreviousSibling() {
		t, isText := n.(*ast.Text)
		if !isText {
			break
		}
		if n != last && (t.SoftLineBreak() || t.HardLineBreak()) {
			break // the final line starts after this node
		}
		run = append([]*ast.Text{t}, run...)
	}
	return run
}