writekit/internal/markdown/codeblock.go

package markdown

import (
	"fmt"
	"html"
	"regexp"
	"strings"

	"github.com/alecthomas/chroma/v2"
	chromahtml "github.com/alecthomas/chroma/v2/formatters/html"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/renderer"
	"github.com/yuin/goldmark/util"
)

var languageIcons = map[string]string{
	"javascript": "javascript",
	"js":         "javascript",
	"typescript": "typescript",
	"ts":         "typescript",
	"html":       "html5",
	"css":        "css3",
	"scss":       "sass",
	"sass":       "sass",
	"less":       "less",
	"react":      "react",
	"jsx":        "react",
	"tsx":        "react",
	"vue":        "vuedotjs",
	"svelte":     "svelte",
	"angular":    "angular",
	"astro":      "astro",
	"python":     "python",
	"py":         "python",
	"go":         "go",
	"golang":     "go",
	"rust":       "rust",
	"java":       "openjdk",
	"kotlin":     "kotlin",
	"scala":      "scala",
	"ruby":       "ruby",
	"rb":         "ruby",
	"php":        "php",
	"csharp":     "csharp",
	"cs":         "csharp",
	"cpp":        "cplusplus",
	"c":          "c",
	"swift":      "swift",
	"json":       "json",
	"yaml":       "yaml",
	"yml":        "yaml",
	"toml":       "toml",
	"xml":        "xml",
	"bash":       "gnubash",
	"sh":         "gnubash",
	"shell":      "gnubash",
	"zsh":        "gnubash",
	"powershell": "powershell",
	"ps1":        "powershell",
	"sql":        "postgresql",
	"mysql":      "mysql",
	"postgres":   "postgresql",
	"postgresql": "postgresql",
	"mongodb":    "mongodb",
	"redis":      "redis",
	"graphql":    "graphql",
	"gql":        "graphql",
	"docker":     "docker",
	"dockerfile": "docker",
	"markdown":   "markdown",
	"md":         "markdown",
	"lua":        "lua",
	"elixir":     "elixir",
	"erlang":     "erlang",
	"haskell":    "haskell",
	"clojure":    "clojure",
	"zig":        "zig",
	"nim":        "nim",
	"r":          "r",
	"julia":      "julia",
	"dart":       "dart",
	"flutter":    "flutter",
	"solidity":   "solidity",
	"terraform":  "terraform",
	"nginx":      "nginx",
	"apache":     "apache",
}

var titleRegex = regexp.MustCompile(`title=["']([^"']+)["']`)

func parseCodeInfo(info string) (language, title string) {
	info = strings.TrimSpace(info)
	if info == "" {
		return "", ""
	}

	if match := titleRegex.FindStringSubmatch(info); match != nil {
		title = match[1]
		info = strings.TrimSpace(titleRegex.ReplaceAllString(info, ""))
	}

	parts := strings.Fields(info)
	if len(parts) > 0 {
		language = parts[0]
	}

	return language, title
}

func getLanguageIconURL(language string) string {
	slug, ok := languageIcons[strings.ToLower(language)]
	if !ok {
		return ""
	}
	return fmt.Sprintf("https://api.iconify.design/simple-icons/%s.svg?color=%%2371717a", slug)
}

type codeBlockRenderer struct {
	style string
}

func (r *codeBlockRenderer) RegisterFuncs(reg renderer.NodeRendererFuncRegisterer) {
	reg.Register(ast.KindFencedCodeBlock, r.renderFencedCodeBlock)
}

func (r *codeBlockRenderer) renderFencedCodeBlock(w util.BufWriter, source []byte, node ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}

	n := node.(*ast.FencedCodeBlock)
	var info string
	if n.Info != nil {
		info = string(n.Info.Segment.Value(source))
	}
	language, title := parseCodeInfo(info)

	var codeContent strings.Builder
	lines := n.Lines()
	for i := 0; i < lines.Len(); i++ {
		line := lines.At(i)
		codeContent.Write(line.Value(source))
	}
	code := codeContent.String()

	hasHeader := title != "" || language != ""

	w.WriteString(`<div class="code-block">`)

	if hasHeader {
		w.WriteString(`<div class="code-header">`)

		if language != "" {
			iconURL := getLanguageIconURL(language)
			if iconURL != "" {
				w.WriteString(fmt.Sprintf(`<img src="%s" alt="%s" class="code-icon" />`, iconURL, html.EscapeString(language)))
			}
		}

		if title != "" {
			w.WriteString(fmt.Sprintf(`<span class="code-title">%s</span>`, html.EscapeString(title)))
		} else if language != "" {
			w.WriteString(fmt.Sprintf(`<span class="code-title">%s</span>`, html.EscapeString(language)))
		}

		w.WriteString(`<button class="code-copy" onclick="navigator.clipboard.writeText(this.closest('.code-block').querySelector('code').textContent)">Copy</button>`)
		w.WriteString(`</div>`)
	}

	highlighted := r.highlightCode(code, language)
	w.WriteString(highlighted)

	w.WriteString(`</div>`)

	return ast.WalkContinue, nil
}

func (r *codeBlockRenderer) highlightCode(code, language string) string {
	lexer := lexers.Get(language)
	if lexer == nil {
		lexer = lexers.Fallback
	}
	lexer = chroma.Coalesce(lexer)

	style := styles.Get(r.style)
	if style == nil {
		style = styles.Fallback
	}

	formatter := chromahtml.New(
		chromahtml.WithClasses(true),
		chromahtml.PreventSurroundingPre(false),
	)

	iterator, err := lexer.Tokenise(nil, code)
	if err != nil {
		return fmt.Sprintf("<pre><code>%s</code></pre>", html.EscapeString(code))
	}

	var buf strings.Builder
	if err := formatter.Format(&buf, style, iterator); err != nil {
		return fmt.Sprintf("<pre><code>%s</code></pre>", html.EscapeString(code))
	}

	return buf.String()
}

type codeBlockExtension struct {
	style string
}

func (e *codeBlockExtension) Extend(m goldmark.Markdown) {
	m.Renderer().AddOptions(
		renderer.WithNodeRenderers(
			util.Prioritized(&codeBlockRenderer{style: e.style}, 100),
		),
	)
}

func NewCodeBlockExtension(style string) goldmark.Extender {
	return &codeBlockExtension{style: style}
}