cthulhu/internal/ocr/engine.go

// Copyright (c) 2026 Storm Dragon
// Authors: Storm Dragon, Angus Kola
// SPDX-License-Identifier: GPL-3.0-or-later

package ocr

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"image"
	"image/png"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	discoveryTimeout   = 2 * time.Second
	recognitionTimeout = 15 * time.Second
	maxCatalogBytes    = 64 << 10
	maxImageBytes      = 64 << 20
	maxStderrBytes     = 16 << 10
)

type Layout int

const (
	AutomaticPage Layout = 3
	SingleBlock   Layout = 6
	SparseText    Layout = 11
)

// Options is a per-job snapshot. Empty language chooses eng when installed,
// otherwise the first installed recognition language in sorted order. An
// explicit language expression is never replaced with a fallback.
type Options struct {
	Language   string
	Layout     Layout // Zero uses AutomaticPage; other unsupported values fail.
	Scale      int    // Zero uses original size; 1 through 4 are explicit factors.
	Grayscale  bool
	Invert     bool
	BlackWhite bool
	Threshold  int
}

// Catalog contains recognition languages, excluding orientation-only osd data.
// Discovery is on demand, not part of reader startup. Each request rescans so
// installing/removing data does not require restarting the reader.
type Catalog struct {
	Languages       []string
	DefaultLanguage string
}

// Engine serializes its subprocesses. Call its blocking methods from a worker,
// never a native callback, GTK thread or the reader's event/input loop. A
// canceled call kills and reaps its process before releasing ownership.
type Engine struct {
	gate    chan struct{}
	command func(context.Context, ...string) *exec.Cmd
}

// NewEngine performs no dependency lookup, native initialization or I/O.
func NewEngine() *Engine {
	return &Engine{gate: make(chan struct{}, 1), command: func(ctx context.Context, args ...string) *exec.Cmd {
		return exec.CommandContext(ctx, "tesseract", args...)
	}}
}

func (engine *Engine) acquire(ctx context.Context) error {
	select {
	case engine.gate <- struct{}{}:
		if err := ctx.Err(); err != nil {
			<-engine.gate
			return err
		}
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (engine *Engine) Discover(ctx context.Context) (Catalog, error) {
	ctx, cancel := context.WithTimeout(ctx, discoveryTimeout)
	defer cancel()
	if err := engine.acquire(ctx); err != nil {
		return Catalog{}, err
	}
	defer func() { <-engine.gate }()
	return engine.discover(ctx)
}

func (engine *Engine) discover(ctx context.Context) (Catalog, error) {
	ctx, cancel := context.WithTimeout(ctx, discoveryTimeout)
	defer cancel()
	output, _, err := engine.run(ctx, nil, maxCatalogBytes, "--list-langs")
	if err != nil {
		if errors.Is(err, ErrRecognition) {
			return Catalog{}, fmt.Errorf("%w: language discovery failed", ErrUnavailable)
		}
		return Catalog{}, err
	}
	return parse_catalog(string(output))
}

func parse_catalog(output string) (Catalog, error) {
	lines := strings.Split(strings.TrimSpace(output), "\n")
	header := strings.TrimSuffix(lines[0], "\r")
	countStart := strings.LastIndex(header, "(")
	if !strings.HasPrefix(header, "List of available languages in ") || !strings.HasSuffix(header, "):") || countStart < 0 {
		return Catalog{}, fmt.Errorf("%w: invalid language catalog", ErrUnavailable)
	}
	count, err := strconv.Atoi(header[countStart+1 : len(header)-2])
	if err != nil || count < 0 || count > 512 || count != len(lines)-1 {
		return Catalog{}, fmt.Errorf("%w: invalid language catalog", ErrUnavailable)
	}
	var catalog Catalog
	seen := make(map[string]bool)
	for _, line := range lines[1:] {
		language := strings.TrimSuffix(line, "\r")
		if !valid_language(language) || seen[language] {
			return Catalog{}, fmt.Errorf("%w: invalid language catalog", ErrUnavailable)
		}
		seen[language] = true
		if language != "osd" {
			catalog.Languages = append(catalog.Languages, language)
		}
	}
	if len(catalog.Languages) == 0 {
		return Catalog{}, fmt.Errorf("%w: no recognition language data installed", ErrUnavailable)
	}
	sort.Strings(catalog.Languages)
	catalog.DefaultLanguage = catalog.Languages[0]
	if seen["eng"] {
		catalog.DefaultLanguage = "eng"
	}
	return catalog, nil
}

// valid_language permits upstream script/Latin-style identifiers, without
// accepting path traversal, options, whitespace or expression delimiters.
func valid_language(language string) bool {
	if language == "" || len(language) > 128 {
		return false
	}
	parts := strings.Split(language, "/")
	if len(parts) > 2 || (len(parts) == 2 && parts[0] != "script") {
		return false
	}
	for _, part := range parts {
		if part == "" || part[0] == '-' {
			return false
		}
		for _, char := range part {
			if !(char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || char >= '0' && char <= '9' || char == '_' || char == '-') {
				return false
			}
		}
	}
	return true
}

func (catalog Catalog) resolve_language(expression string) (string, error) {
	if expression == "" {
		return catalog.DefaultLanguage, nil
	}
	if len(expression) > 1024 {
		return "", fmt.Errorf("%w: invalid language selection", ErrUnavailable)
	}
	seen := make(map[string]bool)
	for _, language := range strings.Split(expression, "+") {
		if !valid_language(language) || language == "osd" || seen[language] {
			return "", fmt.Errorf("%w: invalid language selection", ErrUnavailable)
		}
		seen[language] = true
		index := sort.SearchStrings(catalog.Languages, language)
		if index == len(catalog.Languages) || catalog.Languages[index] != language {
			return "", fmt.Errorf("%w: selected language data is not installed", ErrUnavailable)
		}
	}
	return expression, nil
}

// Recognize encodes one owned image to PNG in memory and obtains text and
// coordinates from one TSV recognition invocation. The caller must not mutate
// pixels until this method returns. No images or text are logged or persisted.
func (engine *Engine) Recognize(ctx context.Context, pixels image.Image, options Options) (Result, error) {
	ctx, cancel := context.WithTimeout(ctx, recognitionTimeout)
	defer cancel()
	if pixels == nil || !valid_dimensions(pixels.Bounds().Dx(), pixels.Bounds().Dy()) {
		return Result{}, ErrInvalidImage
	}
	if options.Layout == 0 {
		options.Layout = AutomaticPage
	}
	if options.Layout != AutomaticPage && options.Layout != SingleBlock && options.Layout != SparseText {
		return Result{}, fmt.Errorf("%w: unsupported text layout", ErrRecognition)
	}
	if err := engine.acquire(ctx); err != nil {
		return Result{}, err
	}
	defer func() { <-engine.gate }()
	var err error
	pixels, err = Preprocess(ctx, pixels, options)
	if err != nil {
		return Result{}, err
	}
	catalog, err := engine.discover(ctx)
	if err != nil {
		return Result{}, err
	}
	language, err := catalog.resolve_language(options.Language)
	if err != nil {
		return Result{}, err
	}
	encoded := &bounded_buffer{limit: maxImageBytes, ctx: ctx}
	if err := png.Encode(encoded, pixels); err != nil {
		if ctx.Err() != nil {
			return Result{}, ctx.Err()
		}
		if errors.Is(err, ErrLimit) {
			return Result{}, ErrLimit
		}
		return Result{}, ErrInvalidImage
	}
	output, diagnostics, err := engine.run(ctx, encoded.Bytes(), maxTSVBytes,
		"stdin", "stdout", "-l", language, "--psm", strconv.Itoa(int(options.Layout)), "tsv")
	// Tesseract may exit successfully after loading only part of a multi-language
	// expression. Never silently accept that degraded selection.
	if strings.Contains(string(diagnostics), "Error opening data file") ||
		strings.Contains(string(diagnostics), "Failed loading language") ||
		strings.Contains(string(diagnostics), "couldn't load any languages") {
		if ctx.Err() != nil {
			return Result{}, ctx.Err()
		}
		return Result{}, fmt.Errorf("%w: selected language data could not be loaded", ErrUnavailable)
	}
	if err != nil {
		return Result{}, err
	}
	if err := ctx.Err(); err != nil {
		return Result{}, err
	}
	result, err := ParseTSV(bytes.NewReader(output), pixels.Bounds().Dx(), pixels.Bounds().Dy())
	if ctx.Err() != nil {
		return Result{}, ctx.Err()
	}
	return result, err
}