khosra/internal/ext/shortcodes/code.go

package shortcodes

import (
	"bytes"
	"io/fs"
	"log/slog"
	"path"
	"strconv"
	"strings"

	"html/template"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/formatters/html"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/parser"
	"github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"

	"khosra/internal/render"
)

// codeFragment is the theme template a highlighted block renders through. The engine produces token spans
// and the theme produces everything around them — the frame, the caption, a copy button (ADR-0075).
const codeFragment = "code"

// codeKind is one fenced block after its info string has been read.
var codeKind = ast.NewNodeKind("ShortcodeCode")

type codeBlock struct {
	ast.BaseBlock
	args map[string]string
	lang string
	// source is what to highlight: the fence's own lines, or the file it named.
	source []byte
	// first is the number the first displayed line carries, and marked are the lines to tint. Both are in the
	// numbering the reader sees, so a slice of a file highlights by that file's line numbers.
	first  int
	marked [][2]int
	// numbered is whether a gutter is drawn at all.
	numbered bool
}

func (n *codeBlock) Kind() ast.NodeKind { return codeKind }

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

// blocks reads every fenced block's info string and replaces the node with one the renderer can highlight.
//
// A transformer, because a block may take its content from a file and only the parse context knows which
// bundle this is — the same reason an include is one (ADR-0038).
type code struct{}

func (code) Transform(doc *ast.Document, reader text.Reader, pc parser.Context) {
	origin, _ := render.OriginFrom(pc)
	source := reader.Source()
	var fenced []*ast.FencedCodeBlock
	_ = ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if entering {
			if f, is := n.(*ast.FencedCodeBlock); is {
				fenced = append(fenced, f)
			}
		}
		return ast.WalkContinue, nil
	})
	for _, f := range fenced {
		if block := readFence(f, source, origin); block != nil {
			f.Parent().ReplaceChild(f.Parent(), f, block)
		}
	}
}

// readFence turns one fenced block into a codeBlock, or nil to leave it exactly as goldmark rendered it.
func readFence(f *ast.FencedCodeBlock, source []byte, origin render.Origin) *codeBlock {
	lang := string(f.Language(source))
	args := map[string]string{}
	if f.Info != nil {
		rest := strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(string(f.Info.Segment.Value(source))), lang))
		for rest != "" {
			key, value, remainder, ok := argument(rest)
			if !ok {
				break
			}
			args[key] = value
			rest = remainder
		}
	}
	block := &codeBlock{args: args, lang: lang, first: 1}
	if name := args["file"]; name != "" {
		body, from, ok := fromFile(origin, name, args["lines"])
		if !ok {
			return nil
		}
		block.source, block.first, block.numbered = body, from, true
	} else {
		var buf bytes.Buffer
		for i := 0; i < f.Lines().Len(); i++ {
			line := f.Lines().At(i)
			buf.Write(line.Value(source))
		}
		block.source = buf.Bytes()
	}
	if start, err := strconv.Atoi(args["start"]); err == nil && start > 0 {
		block.first, block.numbered = start, true
	}
	if args["numbers"] != "" {
		block.numbered = true
	}
	block.marked = ranges(args["hl"])
	return block
}

// fromFile reads a snippet out of a file beside the bundle, and reports the line number it starts at.
//
// The same containment rule an include keeps: a name with `..` is refused, so a block cannot publish a
// template or a dotfile (ADR-0038).
func fromFile(origin render.Origin, name, span string) ([]byte, int, bool) {
	if origin.Files == nil || strings.Contains(name, "..") {
		slog.Error("code block cannot read that file", "file", name)
		return nil, 0, false
	}
	data, err := fs.ReadFile(origin.Files, path.Join(origin.Dir, name))
	if err != nil {
		slog.Error("code block cannot read that file", "file", name, "err", err)
		return nil, 0, false
	}
	lines := strings.Split(strings.TrimRight(string(data), "\n"), "\n")
	from, to := 1, len(lines)
	if pair := ranges(span); len(pair) == 1 {
		from, to = pair[0][0], pair[0][1]
	}
	if from < 1 {
		from = 1
	}
	if to > len(lines) || to < from {
		to = len(lines)
	}
	return []byte(strings.Join(lines[from-1:to], "\n") + "\n"), from, true
}

// ranges reads `3,7-9` into the pairs chroma wants. A number on its own is a range of one.
func ranges(spec string) [][2]int {
	var out [][2]int
	for _, part := range strings.Split(spec, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		lo, hi, split := strings.Cut(part, "-")
		start, err := strconv.Atoi(lo)
		if err != nil {
			continue
		}
		end := start
		if split {
			if end, err = strconv.Atoi(hi); err != nil {
				continue
			}
		}
		out = append(out, [2]int{start, end})
	}
	return out
}

// renderCode highlights the block and hands the result to the theme.
//
// Classes rather than inline colour, so the palette lives in a stylesheet the theme owns and a reader's dark
// mode is the theme's business — the feature decides which token this is and nothing about how it looks
// (ADR-0036, ADR-0075).
func (f fragments) renderCode(w util.BufWriter, source []byte, n ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}
	block := n.(*codeBlock)
	lexer := lexers.Get(block.lang)
	if lexer == nil {
		lexer = lexers.Fallback
	}
	tokens, err := chroma.Coalesce(lexer).Tokenise(nil, string(block.source))
	if err != nil {
		slog.Error("cannot highlight", "lang", block.lang, "err", err)
		return ast.WalkContinue, nil
	}
	options := []html.Option{html.WithClasses(true)}
	if block.numbered {
		options = append(options, html.WithLineNumbers(true), html.BaseLineNumber(block.first))
	}
	if len(block.marked) > 0 {
		options = append(options, html.HighlightLines(block.marked))
	}
	var highlighted bytes.Buffer
	if err := html.New(options...).Format(&highlighted, styles.Fallback, tokens); err != nil {
		slog.Error("cannot format", "lang", block.lang, "err", err)
		return ast.WalkContinue, nil
	}
	args := map[string]string{"lang": block.lang}
	for k, v := range block.args {
		args[k] = v
	}
	out, err := f.partial(codeFragment, render.Fragment{Args: args, Body: template.HTML(highlighted.String())})
	if err != nil {
		slog.Error("skipping code fragment", "err", err)
		out = highlighted.Bytes()
	}
	if _, err := w.Write(out); err != nil {
		return ast.WalkStop, err
	}
	return ast.WalkSkipChildren, nil
}