Implement a few functions

1 files changed, 277 insertions, 24 deletions

diff --git a/main.go b/main.go
index b3a3019..bf63704 100644
--- a/main.go
+++ b/main.go
@@ -1,39 +1,82 @@
 package main
 
 import (
+	"bytes"
+	"crypto/sha1"
 	"database/sql"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"io"
+	"io/fs"
 	"log"
 	"net"
 	"net/http"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"time"
 
 	_ "github.com/mattn/go-sqlite3"
 )
 
-// TODO: Perhaps maintain the DB, and the gallery directory, as globals.
+var (
+	db *sql.DB // sqlite database
+	gd string  // gallery directory
+)
+
+func openDB(directory string) error {
+	var err error
+	db, err = sql.Open("sqlite3", filepath.Join(directory, "gallery.db"))
+	gd = directory
+	return err
+}
 
-func openDB(gd string) (*sql.DB, error) {
-	return sql.Open("sqlite3", filepath.Join(gd, "gallery.db"))
+func imagePath(sha1 string) string {
+	return filepath.Join(gd, "images", sha1[:2], sha1)
 }
 
-// init GD - initialize a "gallery directory" that contains gallery.sqlite,
+func thumbPath(sha1 string) string {
+	return filepath.Join(gd, "thumbs", sha1[:2], sha1+".webp")
+}
+
+func dbCollect(query string) ([]string, error) {
+	rows, err := db.Query(query)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var result []string
+	for rows.Next() {
+		var sha1 string
+		if err := rows.Scan(&sha1); err != nil {
+			return nil, err
+		}
+		result = append(result, sha1)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+// cmdInit initializes a "gallery directory" that contains gallery.sqlite,
 // images, thumbs.
 func cmdInit(args []string) {
 	if len(args) != 1 {
 		log.Fatalln("usage: GD")
 	}
 
-	gd := args[0]
-	db, err := openDB(gd)
-	if err != nil {
+	if err := openDB(args[0]); err != nil {
 		log.Fatalln(err)
 	}
-
-	if _, err = db.Query(initializeSQL); err != nil {
+	if _, err := db.Exec(initializeSQL); err != nil {
 		log.Fatalln(err)
 	}
+
+	// XXX: There's technically no reason to keep images as symlinks,
+	// we might just keep absolute paths in the database as well.
 	if err := os.MkdirAll(filepath.Join(gd, "images"), 0777); err != nil {
 		log.Fatalln(err)
 	}
@@ -42,19 +85,15 @@ func cmdInit(args []string) {
 	}
 }
 
-// run GD ADDRESS - run a web UI against GD on ADDRESS
+// cmdRun runs a web UI against GD on ADDRESS.
 func cmdRun(args []string) {
 	if len(args) != 2 {
 		log.Fatalln("usage: GD ADDRESS")
 	}
-
-	gd := args[0]
-	db, err := openDB(gd)
-	if err != nil {
+	if err := openDB(args[0]); err != nil {
 		log.Fatalln(err)
 	}
 
-	_ = db
 	address := args[1]
 
 	http.Handle("/", http.FileServer(http.Dir("public")))
@@ -83,29 +122,236 @@ func cmdRun(args []string) {
 	log.Fatalln(s.ListenAndServe())
 }
 
-// import GD ROOT... - add files to the "entry" table
+func importFunc(path string, d fs.DirEntry, err error) error {
+	if err != nil || d.IsDir() {
+		return err
+	}
+
+	// The input may be a relative path, and we want to remember it as such,
+	// but symlinks for the images must be absolute.
+	absPath, err := filepath.Abs(path)
+	if err != nil {
+		return err
+	}
+
+	cmd := exec.Command("xdg-mime", "query", "filetype", path)
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+	out, err := io.ReadAll(stdout)
+	if err != nil {
+		return err
+	}
+	if err := cmd.Wait(); err != nil {
+		return err
+	}
+
+	// Skip videos, which ImageMagick can process, but we don't want it to,
+	// so that they're not converted 1:1 to WebP.
+	if !bytes.HasPrefix(out, []byte("image/")) {
+		return nil
+	}
+
+	f, err := os.Open(path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	s, err := f.Stat()
+	if err != nil {
+		return err
+	}
+
+	hash := sha1.New()
+	_, err = io.CopyBuffer(hash, f, make([]byte, 65536))
+	if err != nil {
+		return err
+	}
+
+	hexSHA1 := hex.EncodeToString(hash.Sum(make([]byte, sha1.Size)))
+	// TODO: Make sure that the subdirectory exists.
+	pathImage := imagePath(hexSHA1)
+	if err := os.Symlink(absPath, pathImage); err != nil &&
+		!errors.Is(err, fs.ErrExist) {
+		return err
+	}
+
+	// TODO: This should run in a transaction.
+	dbDirname, dbBasename := filepath.Split(path)
+	_, err = db.Exec(`INSERT INTO entry(
+		path, basename, mtime, sha1
+	) VALUES (?, ?, ?, ?)`, dbDirname, dbBasename, s.ModTime().Unix(), hexSHA1)
+	// TODO: Also ensure that a row exists in "image".
+	return err
+}
+
+// cmdImport adds files to the "entry" table.
 func cmdImport(args []string) {
-	// TODO
+	if len(args) < 1 {
+		log.Fatalln("usage: GD ROOT...")
+	}
+	if err := openDB(args[0]); err != nil {
+		log.Fatalln(err)
+	}
+
+	// TODO: This would better be done in parallel (making hashes).
+	for _, name := range args[1:] {
+		if err := filepath.WalkDir(name, importFunc); err != nil {
+			log.Fatalln(err)
+		}
+	}
 }
 
-// sync GD ROOT... - like import, but clear table beforehands
+// cmdSync is like import, but clears the "entry" table beforehands.
 func cmdSync(args []string) {
+	if len(args) < 1 {
+		log.Fatalln("usage: GD ROOT...")
+	}
+	if err := openDB(args[0]); err != nil {
+		log.Fatalln(err)
+	}
+
 	// TODO
 }
 
-// check GD - see if all files are accessible
+// cmdCheck checks if all files tracked in the DB are accessible.
 func cmdCheck(args []string) {
-	// TODO
+	if len(args) != 1 {
+		log.Fatalln("usage: GD")
+	}
+	if err := openDB(args[0]); err != nil {
+		log.Fatalln(err)
+	}
+
+	// TODO: Check if all hashes of DB entries have a statable image file,
+	// and that all images with thumb{w,h} have a thumbnail file. Perhaps.
 }
 
-// thumbnail GD [SHA1...] - generate missing thumbnails, in parallel
+func makeThumbnail(pathImage, pathThumb string) (int, int, error) {
+	// TODO: Make sure that the trailing subdirectory in pathThumb exists.
+	cmd := exec.Command("convert", pathImage, "-coalesce", "-colorspace", "RGB",
+		"-auto-orient", "-strip", "-resize", "256x128>", "-colorspace", "sRGB",
+		"-format", "%w %h", "+write", "info:", pathThumb)
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return 0, 0, err
+	}
+	if err := cmd.Start(); err != nil {
+		return 0, 0, err
+	}
+	out, err := io.ReadAll(stdout)
+	if err != nil {
+		return 0, 0, err
+	}
+	if err := cmd.Wait(); err != nil {
+		return 0, 0, err
+	}
+
+	var w, h int
+	_, err = fmt.Fscanf(bytes.NewReader(out), "%d %d", &w, &h)
+	return w, h, err
+}
+
+// cmdThumbnail generates missing thumbnails, in parallel.
 func cmdThumbnail(args []string) {
-	// TODO: Show progress.
+	if len(args) < 1 {
+		log.Fatalln("usage: GD [SHA1...]")
+	}
+	if err := openDB(args[0]); err != nil {
+		log.Fatalln(err)
+	}
+
+	hexSHA1 := args[1:]
+	if len(hexSHA1) == 0 {
+		// Get all unique images in the database with no thumbnail.
+		var err error
+		hexSHA1, err = dbCollect(`SELECT DISTINCT sha1 FROM entry
+			LEFT OUTER JOIN image ON entry.sha1 = image.sha1
+			WHERE thumbw IS NULL OR thumbh IS NULL`)
+		if err != nil {
+			log.Fatalln(err)
+		}
+	}
+
+	// TODO: Try to run the thumbnailer in parallel, somehow.
+	// TODO: Show progress in some manner. Perhaps port my propeller code.
+	for _, sha1 := range hexSHA1 {
+		pathImage := imagePath(sha1)
+		pathThumb := thumbPath(sha1)
+		w, h, err := makeThumbnail(pathImage, pathThumb)
+		if err != nil {
+			log.Fatalln(err)
+		}
+
+		_, err = db.Exec(`INSERT INTO image(
+			sha1, thumbw, thumbh, dhash
+		) VALUES (?, ?, ?, NULL)`, sha1, w, h)
+		if err != nil {
+			log.Fatalln(err)
+		}
+	}
+}
+
+func makeDhash(hasher, pathThumb string) (uint64, error) {
+	cmd := exec.Command(hasher, pathThumb)
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return 0, err
+	}
+	if err := cmd.Start(); err != nil {
+		return 0, err
+	}
+	out, err := io.ReadAll(stdout)
+	if err != nil {
+		return 0, err
+	}
+	if err := cmd.Wait(); err != nil {
+		return 0, err
+	}
+
+	var hash uint64
+	_, err = fmt.Fscanf(bytes.NewReader(out), "%d", &hash)
+	return hash, err
 }
 
-// dhash GD HASHER [SHA1...] - generate perceptual hash from thumbnails
+// cmdDhash generates perceptual hash from thumbnails.
 func cmdDhash(args []string) {
-	// TODO
+	if len(args) < 1 {
+		log.Fatalln("usage: GD HASHER [SHA1...]")
+	}
+	if err := openDB(args[0]); err != nil {
+		log.Fatalln(err)
+	}
+
+	hasher, hexSHA1 := args[1], args[2:]
+	if len(hexSHA1) == 0 {
+		var err error
+		hexSHA1, err = dbCollect(`SELECT sha1 FROM image WHERE dhash IS NULL`)
+		if err != nil {
+			log.Fatalln(err)
+		}
+	}
+
+	// TODO: Try to run the hasher in parallel, somehow.
+	for _, sha1 := range hexSHA1 {
+		pathThumb := thumbPath(sha1)
+		hash, err := makeDhash(hasher, pathThumb)
+		if err != nil {
+			log.Fatalln(err)
+		}
+
+		_, err = db.Exec(`UPDATE image SET dhash = ? WHERE sha1 = ?`,
+			int64(hash), sha1)
+		if err != nil {
+			log.Fatalln(err)
+		}
+	}
 }
 
 var commands = map[string]struct {
@@ -130,5 +376,12 @@ func main() {
 		log.Fatalln("Unknown command: " + os.Args[1])
 	}
 
+	// TODO: Check if this runs on fatal errors.
+	defer func() {
+		if db != nil {
+			db.Close()
+		}
+	}()
+
 	cmd.handler(os.Args[2:])
 }