orientator/internal/pipeline/pipeline.go

// Package pipeline orchestrates orientation detection and rotation across
// images, PDFs, and archives.
package pipeline

import (
	"archive/zip"
	"bytes"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"path/filepath"
	"strings"
	"sync"

	_ "golang.org/x/image/tiff"
	_ "golang.org/x/image/webp"

	"orientator/internal/detector"
	"orientator/internal/extract"
	"orientator/internal/rasterize"
)

// PageReport describes what happened to a single page or image.
type PageReport struct {
	File       string  `json:"file"`
	Page       int     `json:"page,omitempty"` // 0 for plain images
	Angle      int     `json:"angle"`
	Confidence float64 `json:"confidence"`
}

// Result holds the rewritten output and a per-page report.
type Result struct {
	OutputName  string
	OutputBytes []byte
	Reports     []PageReport
}

// Process processes a single uploaded file and returns the corrected output.
// For archives the output is a zip of all corrected entries.
func Process(name string, data []byte) (*Result, error) {
	det, err := detector.Get()
	if err != nil {
		return nil, fmt.Errorf("detector init: %w", err)
	}

	if extract.IsArchive(name) {
		return processArchive(det, name, data)
	}
	return processSingle(det, name, data)
}

func processSingle(det *detector.Detector, name string, data []byte) (*Result, error) {
	ext := strings.ToLower(filepath.Ext(name))
	switch ext {
	case ".pdf":
		out, reports, err := processPDF(det, name, data)
		if err != nil {
			return nil, err
		}
		return &Result{OutputName: name, OutputBytes: out, Reports: reports}, nil
	case ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp":
		out, outName, report, err := processImage(det, name, data)
		if err != nil {
			return nil, err
		}
		return &Result{OutputName: outName, OutputBytes: out, Reports: []PageReport{report}}, nil
	}
	return nil, fmt.Errorf("unsupported file type: %s", ext)
}

func processImage(det *detector.Detector, name string, data []byte) ([]byte, string, PageReport, error) {
	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, "", PageReport{}, fmt.Errorf("decode %s: %w", name, err)
	}
	angle, conf, err := det.Predict(img)
	if err != nil {
		return nil, "", PageReport{}, err
	}
	rotated := RotateImage(img, angle)
	out, ext, err := EncodeImage(rotated, format)
	if err != nil {
		return nil, "", PageReport{}, err
	}
	outName := replaceExt(name, "."+ext)
	return out, outName, PageReport{File: name, Angle: angle, Confidence: conf}, nil
}

func processPDF(det *detector.Detector, name string, data []byte) ([]byte, []PageReport, error) {
	pdf, err := rasterize.Open(data)
	if err != nil {
		return nil, nil, err
	}
	defer pdf.Close()

	n := pdf.NumPages()
	angles := make([]int, n)
	reports := make([]PageReport, 0, n)
	for i := 0; i < n; i++ {
		img, err := pdf.RenderPage(i)
		if err != nil {
			return nil, nil, err
		}
		angle, conf, err := det.Predict(img)
		if err != nil {
			return nil, nil, err
		}
		angles[i] = angle
		reports = append(reports, PageReport{File: name, Page: i + 1, Angle: angle, Confidence: conf})
	}
	out, err := rasterize.RotatePDFPages(data, angles)
	if err != nil {
		return nil, nil, err
	}
	return out, reports, nil
}

func processArchive(det *detector.Detector, name string, data []byte) (*Result, error) {
	var (
		mu      sync.Mutex
		buf     bytes.Buffer
		zw      = zip.NewWriter(&buf)
		reports []PageReport
	)

	err := extract.Walk(name, data, func(e extract.Entry) error {
		// Recurse into nested archives.
		if extract.IsArchive(e.Name) {
			sub, err := Process(e.Name, e.Data)
			if err != nil {
				return writeZipFailure(zw, e.Name, err)
			}
			mu.Lock()
			reports = append(reports, sub.Reports...)
			mu.Unlock()
			return writeZipEntry(zw, sub.OutputName, sub.OutputBytes)
		}

		ext := strings.ToLower(filepath.Ext(e.Name))
		switch ext {
		case ".pdf":
			out, rs, err := processPDF(det, e.Name, e.Data)
			if err != nil {
				return writeZipFailure(zw, e.Name, err)
			}
			mu.Lock()
			reports = append(reports, rs...)
			mu.Unlock()
			return writeZipEntry(zw, e.Name, out)
		case ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp":
			out, outName, r, err := processImage(det, e.Name, e.Data)
			if err != nil {
				return writeZipFailure(zw, e.Name, err)
			}
			mu.Lock()
			reports = append(reports, r)
			mu.Unlock()
			return writeZipEntry(zw, outName, out)
		default:
			// Pass-through any non-document files.
			return writeZipEntry(zw, e.Name, e.Data)
		}
	})
	if err != nil {
		return nil, err
	}
	if err := zw.Close(); err != nil {
		return nil, err
	}

	outName := replaceExt(name, ".zip")
	return &Result{OutputName: outName, OutputBytes: buf.Bytes(), Reports: reports}, nil
}

func writeZipEntry(zw *zip.Writer, name string, data []byte) error {
	w, err := zw.Create(name)
	if err != nil {
		return err
	}
	_, err = io.Copy(w, bytes.NewReader(data))
	return err
}

func writeZipFailure(zw *zip.Writer, name string, cause error) error {
	w, err := zw.Create(name + ".error.txt")
	if err != nil {
		return err
	}
	_, _ = w.Write([]byte(cause.Error()))
	return nil
}

func replaceExt(name, newExt string) string {
	old := filepath.Ext(name)
	return strings.TrimSuffix(name, old) + newExt
}