package render

import (
	"html"

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

// Mermaid support: a fenced ```mermaid code block is rendered as
// <pre class="mermaid">…diagram source…</pre> instead of being syntax
// highlighted, so the client-side mermaid script (loaded by the page template
// only when such an element is present) can draw it. The diagram source is
// HTML-escaped here and passes through the bluemonday sanitizer as plain text;
// mermaid reads it back via the element's textContent.

// kindMermaid is the AST node kind for a mermaid diagram block.
var kindMermaid = ast.NewNodeKind("Mermaid")

type mermaidNode struct {
	ast.BaseBlock
	source []byte
}

func (n *mermaidNode) Kind() ast.NodeKind { return kindMermaid }

func (n *mermaidNode) Dump(source []byte, level int) {
	ast.DumpHelper(n, source, level, nil, nil)
}

// mermaidTransformer rewrites ```mermaid fenced code blocks into mermaidNode,
// before the highlighting renderer would otherwise turn them into code.
type mermaidTransformer struct{}

func (mermaidTransformer) Transform(doc *ast.Document, reader text.Reader, _ parser.Context) {
	source := reader.Source()
	var found []*ast.FencedCodeBlock
	_ = ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		if fc, ok := n.(*ast.FencedCodeBlock); ok && string(fc.Language(source)) == "mermaid" {
			found = append(found, fc)
		}
		return ast.WalkContinue, nil
	})
	for _, fc := range found {
		var buf []byte
		for i := 0; i < fc.Lines().Len(); i++ {
			seg := fc.Lines().At(i)
			buf = append(buf, seg.Value(source)...)
		}
		m := &mermaidNode{source: buf}
		if p := fc.Parent(); p != nil {
			p.ReplaceChild(p, fc, m)
		}
	}
}

// mermaidHTMLRenderer renders a mermaidNode.
type mermaidHTMLRenderer struct{}

func (mermaidHTMLRenderer) RegisterFuncs(reg renderer.NodeRendererFuncRegisterer) {
	reg.Register(kindMermaid, renderMermaid)
}

func renderMermaid(w util.BufWriter, _ []byte, n ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}
	m := n.(*mermaidNode)
	_, _ = w.WriteString(`<pre class="mermaid">`)
	_, _ = w.WriteString(html.EscapeString(string(m.source)))
	_, _ = w.WriteString("</pre>\n")
	return ast.WalkSkipChildren, nil
}