markdown.go

package main

import (
	"bytes"
	"fmt"
	"regexp"
	"strconv"

	"github.com/yuin/goldmark"
	highlighting "github.com/yuin/goldmark-highlighting/v2"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/extension"
	"github.com/yuin/goldmark/parser"
	gmhtml "github.com/yuin/goldmark/renderer/html"
	"github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"
)

// codeBlockInfo stores line numbers for a code block
type codeBlockInfo struct {
	startLine int
	endLine   int
}

// LineAttributeTransformer collects line numbers for all block nodes
type LineAttributeTransformer struct {
	codeBlocks []*codeBlockInfo
}

func (t *LineAttributeTransformer) Transform(doc *ast.Document, reader text.Reader, pc parser.Context) {
	source := reader.Source()
	t.codeBlocks = nil // Reset for each document

	_ = ast.Walk(doc, func(node ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}

		// Only process block-level nodes
		if node.Type() == ast.TypeBlock {
			// Skip list containers (ul/ol) - their children (li) will have attributes
			if node.Kind() == ast.KindList {
				return ast.WalkContinue, nil
			}

			var startLine, endLine int

			// Special handling for FencedCodeBlock to include the opening fence line
			if node.Kind() == ast.KindFencedCodeBlock {
				fcb := node.(*ast.FencedCodeBlock)
				// Use the Info segment to find the opening fence line
				if fcb.Info != nil {
					infoStart := fcb.Info.Segment.Start
					startLine = bytes.Count(source[:infoStart], []byte{'\n'}) + 1
				} else {
					// No info, use first line of content
					if fcb.Lines().Len() > 0 {
						firstLine := fcb.Lines().At(0)
						// The opening fence is on the line before the first content line
						startLine = bytes.Count(source[:firstLine.Start], []byte{'\n'})
					}
				}

				// End line is after the last content line (includes closing fence)
				if fcb.Lines().Len() > 0 {
					lastLine := fcb.Lines().At(fcb.Lines().Len() - 1)
					endLine = bytes.Count(source[:lastLine.Stop], []byte{'\n'}) + 1
					// Add 1 for the closing fence line
					endLine++
				}

				// Store code block info for post-processing
				t.codeBlocks = append(t.codeBlocks, &codeBlockInfo{
					startLine: startLine,
					endLine:   endLine,
				})
			} else {
				lines := node.Lines()

				if lines.Len() > 0 {
					// Node has direct line info
					firstLine := lines.At(0)
					startLine = bytes.Count(source[:firstLine.Start], []byte{'\n'}) + 1

					lastLine := lines.At(lines.Len() - 1)
					endLine = bytes.Count(source[:lastLine.Stop], []byte{'\n'}) + 1
				} else {
					// Node has no direct line info
					// Calculate from children
					startLine, endLine = getChildLineRange(node, source)
					if startLine == 0 {
						// No line info available from children either
						return ast.WalkContinue, nil
					}
				}
			}

			// Set attributes (goldmark's HTML renderer will automatically render them for non-code blocks)
			node.SetAttribute([]byte("data-line-start"), []byte(strconv.Itoa(startLine)))
			node.SetAttribute([]byte("data-line-end"), []byte(strconv.Itoa(endLine)))
		}

		return ast.WalkContinue, nil
	})
}

// getChildLineRange calculates line range from a node's children
func getChildLineRange(node ast.Node, source []byte) (int, int) {
	var startLine, endLine int

	// Walk children to find first and last line numbers
	for child := node.FirstChild(); child != nil; child = child.NextSibling() {
		lines := child.Lines()
		if lines.Len() > 0 {
			firstLine := lines.At(0)
			childStart := bytes.Count(source[:firstLine.Start], []byte{'\n'}) + 1

			lastLine := lines.At(lines.Len() - 1)
			childEnd := bytes.Count(source[:lastLine.Stop], []byte{'\n'}) + 1

			if startLine == 0 || childStart < startLine {
				startLine = childStart
			}
			if childEnd > endLine {
				endLine = childEnd
			}
		} else {
			// Recursively check grandchildren
			childStart, childEnd := getChildLineRange(child, source)
			if childStart > 0 {
				if startLine == 0 || childStart < startLine {
					startLine = childStart
				}
				if childEnd > endLine {
					endLine = childEnd
				}
			}
		}
	}

	return startLine, endLine
}

// LineAttributeExtension is a goldmark extension that adds line number attributes
type LineAttributeExtension struct {
	transformer *LineAttributeTransformer
}

func (e *LineAttributeExtension) Extend(m goldmark.Markdown) {
	m.Parser().AddOptions(
		parser.WithASTTransformers(
			util.Prioritized(e.transformer, 100),
		),
	)
}

// addLineAttributesToCodeBlocks post-processes HTML to add data-line-* attributes to <pre> tags
// generated by Chroma. This allows us to use Chroma's default rendering while still adding
// our custom line tracking attributes.
func addLineAttributesToCodeBlocks(html []byte, codeBlocks []*codeBlockInfo) []byte {
	// Match <pre> tags (may have style attribute from Chroma)
	re := regexp.MustCompile(`<pre([^>]*)>`)
	blockIndex := 0

	result := re.ReplaceAllFunc(html, func(match []byte) []byte {
		if blockIndex >= len(codeBlocks) {
			// Not a code block (or we've run out of tracked blocks), return as-is
			return match
		}

		block := codeBlocks[blockIndex]
		blockIndex++

		// Extract existing attributes (e.g., style="background-color:#f0f0f0;")
		existingAttrs := string(match[4 : len(match)-1])

		// Add our custom attributes at the beginning
		return []byte(fmt.Sprintf(`<pre data-line-start="%d" data-line-end="%d"%s>`,
			block.startLine, block.endLine, existingAttrs))
	})

	return result
}

// RenderMarkdownWithLineNumbers renders markdown to HTML with line number attributes
func RenderMarkdownWithLineNumbers(source []byte) ([]byte, error) {
	transformer := &LineAttributeTransformer{}

	md := goldmark.New(
		goldmark.WithExtensions(
			extension.GFM,
			&LineAttributeExtension{transformer: transformer},
			highlighting.NewHighlighting(
				highlighting.WithStyle("friendly"),
			),
		),
		goldmark.WithRendererOptions(
			gmhtml.WithUnsafe(), // Allow raw HTML
		),
	)

	var buf bytes.Buffer
	if err := md.Convert(source, &buf); err != nil {
		return nil, err
	}

	// Post-process HTML to add data-line-* attributes to code blocks
	html := addLineAttributesToCodeBlocks(buf.Bytes(), transformer.codeBlocks)

	return html, nil
}

// RenderMarkdown renders markdown to HTML without line number attributes
func RenderMarkdown(source []byte) ([]byte, error) {
	md := goldmark.New(
		goldmark.WithExtensions(
			extension.GFM,
			highlighting.NewHighlighting(
				highlighting.WithStyle("friendly"),
			),
		),
		goldmark.WithRendererOptions(
			gmhtml.WithUnsafe(), // Allow raw HTML
		),
	)

	var buf bytes.Buffer
	if err := md.Convert(source, &buf); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}