oddmu/static_cmd.go

package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"image/jpeg"
	"io"
	"io/fs"
	"net/url"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"github.com/edwvee/exiffix"
	"github.com/gen2brain/webp"
	"github.com/gomarkdown/markdown"
	"github.com/gomarkdown/markdown/ast"
	"github.com/gomarkdown/markdown/html"
	"github.com/google/subcommands"
)

var shrinkWidth = 800
var shrinkQuality = 30

type staticCmd struct {
	jobs    int
	shrink  bool
	glob    string
	verbose bool
}

func (cmd *staticCmd) SetFlags(f *flag.FlagSet) {
	f.IntVar(&cmd.jobs, "jobs", 2, "how many jobs to use")
	f.BoolVar(&cmd.shrink, "shrink", false, "shrink images by decreasing the quality")
	f.StringVar(&cmd.glob, "glob", "", "only export files matching this shell file name pattern")
	f.BoolVar(&cmd.verbose, "verbose", false, "print the files as they are being processed")
}

func (*staticCmd) Name() string     { return "static" }
func (*staticCmd) Synopsis() string { return "generate static HTML files for all pages" }
func (*staticCmd) Usage() string {
	return `static [-jobs n] <dir name>:
  Create static copies in the given directory. Per default, two jobs
  are used to read and write files, but more can be assigned.
`
}

func (cmd *staticCmd) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	args := f.Args()
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "Exactly one target directory is require", args)
		return subcommands.ExitFailure
	}
	dir := filepath.Clean(args[0])
	return staticCli(".", dir, cmd.jobs, cmd.glob, cmd.shrink, cmd.verbose, false)
}

type args struct {
	source, target string
	info           fs.FileInfo
}

// staticCli generates a static site in the designated directory. The quiet flag is used to suppress output when running
// tests. The source directory cannot be set from the command-line. The current directory (".") is assumed.
func staticCli(source, target string, jobs int, glob string, shrink, verbose, quiet bool) subcommands.ExitStatus {
	index.load()
	index.RLock()
	defer index.RUnlock()
	loadLanguages()
	initTemplates()
	tasks := make(chan args, 10000)
	results := make(chan error, jobs)
	done := make(chan bool, jobs)
	stop := make(chan error)
	for i := 0; i < jobs; i++ {
		go staticWorker(tasks, results, done, shrink, verbose)
	}
	go staticWalk(source, target, glob, tasks, stop)
	go staticWatch(jobs, results, done)
	n, err := staticProgressIndicator(results, stop, quiet)
	if !quiet {
		fmt.Printf("\r%d files processed\n", n)
	}
	if err != nil {
		fmt.Println(err)
		return subcommands.ExitFailure
	}
	return subcommands.ExitSuccess
}

// staticWalk walks the source directory tree. Any directory it finds, it recreates in the target directory. Any file it
// finds, it puts into the tasks channel for the staticWorker. When the directory walk is finished, the tasks channel is
// closed. If there's an error on the stop channel, the walk returns that error.
func staticWalk(source, target, glob string, tasks chan (args), stop chan (error)) {
	// The error returned here is what's in the stop channel but at the very end, a worker might return an error
	// even though the walk is already done. This is why we cannot rely on the return value of the walk.
	n := 0
	err := filepath.Walk(source, func(fp string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// don't wait for the stop channel
		select {
		case err := <-stop:
			return err
		default:
			// skip hidden directories and files
			if fp != "." && strings.HasPrefix(filepath.Base(fp), ".") {
				if info.IsDir() {
					return filepath.SkipDir
				}
				return nil
			}
			// skip backup files, avoid recursion
			if strings.HasSuffix(fp, "~") || strings.HasPrefix(fp, target) {
				return nil
			}
			// skip templates
			if slices.Contains(templateFiles, filepath.Base(fp)) {
				return nil
			}
			// skip files that don't match the glob, if set
			if fp != "." && glob != "" {
				match, err := filepath.Match(glob, fp)
				if err != nil {
					return err // abort
				}
				if !match {
					if info.IsDir() {
						return filepath.SkipDir
					}
					return nil
				}
			}
			// determine the actual target: if source is a/ and target is b/ and path is a/file, then the
			// target is b/file
			var actualTarget string
			if source == "." {
				actualTarget = filepath.Join(target, fp)
			} else {
				if !strings.HasPrefix(fp, source) {
					return fmt.Errorf("%s is not a subdirectory of %s", fp, source)
				}
				actualTarget = filepath.Join(target, fp[len(source):])
			}
			// recreate subdirectories, ignore existing ones
			if info.IsDir() {
				os.Mkdir(actualTarget, 0755)
				return nil
			}
			// Markdown files end up as HTML files
			if strings.HasSuffix(actualTarget, ".md") {
				actualTarget = actualTarget[:len(actualTarget)-3] + ".html"
			}
			// do the task if the target file doesn't exist or if the source file is newer
			other, err := os.Stat(actualTarget)
			if err != nil || info.ModTime().After(other.ModTime()) {
				if err == nil {
					fmt.Println(fp, info.ModTime(), other.ModTime(), info.ModTime().After(other.ModTime()))
				}
				n++
				tasks <- args{source: fp, target: actualTarget, info: info}
			}
			return nil
		}
	})
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Printf("\r%d files to process\n", n)
	}
	close(tasks)
}

// staticWatch counts the values coming out of the done channel. When the count matches the number of jobs started, we
// know that all the tasks have been processed and the results channel is closed.
func staticWatch(jobs int, results chan (error), done chan (bool)) {
	for i := 0; i < jobs; i++ {
		<-done
	}
	close(results)
}

// staticWorker takes arguments off the tasks channel (the file to process) and put results in the results channel (any
// errors encountered); when they're done they send true on the done channel.
func staticWorker(tasks chan (args), results chan (error), done chan (bool), shrink, verbose bool) {
	task, ok := <-tasks
	for ok {
		if verbose {
			fmt.Println(task.source)
		}
		results <- staticFile(task.source, task.target, task.info, shrink)
		task, ok = <-tasks
	}
	done <- true
}

// staticProgressIndicator watches the results channel and prints a running count. If the result channel reports an
// error, that is put into the stop channel so that staticWalk stops adding to the tasks channel.
func staticProgressIndicator(results chan (error), stop chan (error), quiet bool) (int, error) {
	n := 0
	t := time.Now()
	err, ok := <-results
	for ok && err == nil {
		n++
		if !quiet && n%13 == 0 {
			if time.Since(t) > time.Second {
				fmt.Printf("\r%d", n)
				t = time.Now()
			}
		}
		err, ok = <-results
	}
	if ok && err != nil {
		// this stops the walker from adding more tasks
		stop <- err
	}
	return n, err
}

// staticFile is used to walk the file trees and do the right thing for the destination directory: create
// subdirectories, link files, render HTML files.
func staticFile(source, target string, info fs.FileInfo, shrink bool) error {
	// render pages
	if strings.HasSuffix(source, ".md") {
		// target already has ".html" extension
		p, err := staticPage(source[:len(source)-3], target)
		if err != nil {
			return err
		}
		return staticFeed(source[:len(source)-3], target[:len(target)-5]+".rss", p, info.ModTime())
	}
	if shrink {
		switch filepath.Ext(source) {
		case ".jpg", ".jpeg", ".webp":
			return shrinkImage(source, target, info)
		}
	}
	// delete before linking, ignore errors
	os.Remove(target)
	err := os.Link(source, target)
	if err == nil {
		return nil
	}
	// in case of invalid cross-device link error, copy file instead
	src, err := os.Open(source)
	if err != nil {
		return err
	}
	defer src.Close()
	dst, err := os.Create(target)
	if err != nil {
		return err
	}
	defer dst.Close()
	_, err = io.Copy(dst, src)
	return err
}

// shrink Image shrinks images down and reduces the quality dramatically.
func shrinkImage(source, target string, info fs.FileInfo) error {
	file, err := os.Open(source)
	if err != nil {
		return err
	}
	defer file.Close()
	img, _, err := exiffix.Decode(file)
	if err != nil {
		return fmt.Errorf("%s cannot be decoded", source)
	}
	if img.Bounds().Dx() > shrinkWidth {
		res := imaging.Resize(img, shrinkWidth, 0, imaging.Lanczos) // preserve aspect ratio
		// imaging functions don't return errors but empty images…
		if res.Rect.Empty() {
			return fmt.Errorf("%s cannot be resized", source)
		}
		img = res
	}
	dst, err := os.Create(target)
	if err != nil {
		return err
	}
	defer dst.Close()
	switch filepath.Ext(source) {
	case ".jpg", ".jpeg":
		err = jpeg.Encode(dst, img, &jpeg.Options{Quality: shrinkQuality})
	case ".webp":
		err = webp.Encode(dst, img, webp.Options{Quality: shrinkQuality})
	}
	return err
}

// staticPage takes the filename of a page (ending in ".md") and generates a static HTML page.
func staticPage(source, target string) (*Page, error) {
	p, err := loadPage(filepath.ToSlash(source))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot load %s: %s\n", source, err)
		return nil, err
	}
	p.handleTitle(true)
	// instead of p.renderHTML() we do it all ourselves, appending ".html" to all the local links
	parser, hashtags := wikiParser()
	doc := markdown.Parse(p.Body, parser)
	ast.WalkFunc(doc, staticLinks)
	opts := html.RendererOptions{
		// sync with wikiRenderer
		Flags: html.CommonFlags & ^html.SmartypantsFractions | html.LazyLoadImages,
	}
	renderer := html.NewRenderer(opts)
	maybeUnsafeHTML := markdown.Render(doc, renderer)
	p.HTML = unsafeBytes(maybeUnsafeHTML)
	p.Hashtags = *hashtags
	return p, write(p, target, "", "static.html")
}

// staticFeed writes a .rss file for a page, but only if it's an index page or a page that might be used as a hashtag
func staticFeed(source, target string, p *Page, ti time.Time) error {
	// render feed, maybe
	base := filepath.Base(source)
	_, ok := index.token[strings.ToLower(base)]
	if base == "index" || ok {
		f := feed(p, ti, 0, 10, ModTime)
		if len(f.Items) > 0 {
			return write(f, target, `<?xml version="1.0" encoding="UTF-8"?>`, "feed.html")
		}
	}
	return nil
}

// staticLinks checks a node and if it is a link to a local page, it appends ".html" to the link destination.
func staticLinks(node ast.Node, entering bool) ast.WalkStatus {
	if entering {
		if v, ok := node.(*ast.Link); ok {
			// not an absolute URL, not a full URL, not a mailto: URI
			if !bytes.HasPrefix(v.Destination, []byte("/")) &&
				!bytes.Contains(v.Destination, []byte("://")) &&
				!bytes.HasPrefix(v.Destination, []byte("mailto:")) {
				// pointing to a page file (instead of an image file, for example).
				fn, err := url.PathUnescape(string(v.Destination))
				if err != nil {
					return ast.GoToNext
				}
				_, err = os.Stat(fn + ".md")
				if err != nil {
					return ast.GoToNext
				}
				v.Destination = append(v.Destination, []byte(".html")...)
			}
		}
	}
	return ast.GoToNext
}

// write a page or feed with an appropriate template to a specific destination, overwriting it.
func write(data any, fp, prefix, templateFile string) error {
	file, err := os.Create(fp)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot create %s: %s\n", fp, err)
		return err
	}
	_, err = file.Write([]byte(prefix))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot write prefix %s: %s\n", fp, err)
		return err
	}
	templates.RLock()
	defer templates.RUnlock()
	err = templates.template[templateFile].Execute(file, data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot execute %s template for %s: %s\n", templateFile, fp, err)
		return err
	}
	return nil
}