ragflow/internal/deepdoc/parser/pdf/pdfium/pdfium.go

//go:build cgo

// Package pdfium renders PDF pages using libpdfium (statically linked
// at build time via CGO_LDFLAGS). It exists solely to replace pdf_oxide's
// RenderPageRaw for use cases where image quality matters for downstream
// OCR/DLA — pdf_oxide still handles all text/char/table extraction.
package pdfium

/*
#cgo LDFLAGS: -lm -lpthread -ldl

#include <stdint.h>
#include <stdlib.h>

typedef struct FPDF_DOCUMENT__  { int unused; } *FPDF_DOCUMENT;
typedef struct FPDF_PAGE__     { int unused; } *FPDF_PAGE;
typedef struct FPDF_BITMAP__   { int unused; } *FPDF_BITMAP;
typedef struct FPDF_BOOKMARK__ { int unused; } *FPDF_BOOKMARK;
typedef struct FPDF_DEST__     { int unused; } *FPDF_DEST;

extern void          FPDF_InitLibrary(void);
extern FPDF_DOCUMENT FPDF_LoadMemDocument(const void* data_buf, int size, const char* password);
extern void          FPDF_CloseDocument(FPDF_DOCUMENT document);
extern int           FPDF_GetPageCount(FPDF_DOCUMENT document);
extern FPDF_PAGE     FPDF_LoadPage(FPDF_DOCUMENT document, int page_index);
extern void          FPDF_ClosePage(FPDF_PAGE page);
extern double        FPDF_GetPageWidth(FPDF_PAGE page);
extern double        FPDF_GetPageHeight(FPDF_PAGE page);
extern FPDF_BITMAP   FPDFBitmap_Create(int width, int height, int alpha);
extern void          FPDFBitmap_Destroy(FPDF_BITMAP bitmap);
extern void          FPDF_RenderPageBitmap(FPDF_BITMAP bitmap, FPDF_PAGE page,
                       int start_x, int start_y, int size_x, int size_y,
                       int rotate, int flags);
extern void*         FPDFBitmap_GetBuffer(FPDF_BITMAP bitmap);
extern int           FPDFBitmap_GetWidth(FPDF_BITMAP bitmap);
extern int           FPDFBitmap_GetHeight(FPDF_BITMAP bitmap);
extern int           FPDFBitmap_GetStride(FPDF_BITMAP bitmap);

// Outline / bookmark API
extern FPDF_BOOKMARK FPDFBookmark_GetFirstChild(FPDF_DOCUMENT document, FPDF_BOOKMARK bookmark);
extern FPDF_BOOKMARK FPDFBookmark_GetNextSibling(FPDF_DOCUMENT document, FPDF_BOOKMARK bookmark);
extern unsigned long  FPDFBookmark_GetTitle(FPDF_BOOKMARK bookmark, void* buffer, unsigned long buflen);
extern FPDF_DEST      FPDFBookmark_GetDest(FPDF_DOCUMENT document, FPDF_BOOKMARK bookmark);
extern int            FPDFDest_GetDestPageIndex(FPDF_DOCUMENT document, FPDF_DEST dest);
*/
import "C"
import (
	"fmt"
	"image"
	"image/color"
	"math"
	"sync"
	"unicode/utf16"
	"unsafe"

	pdfsync "ragflow/internal/deepdoc/parser/pdf/pdfsync"
)

// Outline represents one entry in a PDF document outline (table of contents).
type Outline struct {
	Title      string
	Level      int
	PageNumber int // 1-indexed, matching Python
}

var initOnce sync.Once

// Init initializes the PDFium library. Safe to call multiple times.
func Init() { initOnce.Do(func() { C.FPDF_InitLibrary() }) }

// PageSize returns the page dimensions in PDF points (1/72 inch) as seen
// after rotation.  For a page with /Rotate 90 on A4, this returns ~842×595
// (swapped from the MediaBox 595×842).  The call is cheap — it opens the
// document and page, reads dimensions, then closes.
func PageSize(pdfData []byte, pageIdx int) (width, height float64, err error) {
	pdfsync.With(func() {
		Init()
		_, _, pw, ph, closeAll, e := openPage(pdfData, pageIdx)
		if e != nil {
			err = e
			return
		}
		closeAll()
		width, height = pw, ph
	})
	return
}

// RenderPage renders a single page of a PDF to an *image.RGBA at the given DPI.
// pdfData is the raw PDF bytes, pageIdx is 0-based.
func RenderPage(pdfData []byte, pageIdx int, dpi float64) (out *image.RGBA, err error) {
	pdfsync.With(func() {
		Init()
		_, page, pw, ph, closeAll, e := openPage(pdfData, pageIdx)
		if e != nil {
			err = e
			return
		}
		defer closeAll()

		scale := dpi / 72.0
		pxW := int(math.Round(pw * scale))
		pxH := int(math.Round(ph * scale))

		bitmap := C.FPDFBitmap_Create(C.int(pxW), C.int(pxH), 1) // 1 = RGBA
		if bitmap == nil {
			err = fmt.Errorf("pdfium: FPDFBitmap_Create(%d,%d) returned nil", pxW, pxH)
			return
		}
		defer C.FPDFBitmap_Destroy(bitmap)

		// Fill with opaque white before rendering, so transparent areas
		// (e.g. outside crop box) are white rather than undefined.
		stride := int(C.FPDFBitmap_GetStride(bitmap))
		buf := C.FPDFBitmap_GetBuffer(bitmap)
		pixels := (*[1 << 30]byte)(unsafe.Pointer(buf))[: pxH*stride : pxH*stride]
		for i := range pixels {
			pixels[i] = 255
		}

		// FPDF_ANNOT (0x01) — render annotations.
		// LCD text AA (0x02) is left off; default text smoothing is sufficient.
		C.FPDF_RenderPageBitmap(bitmap, page, 0, 0, C.int(pxW), C.int(pxH), 0, 0x01)

		// pdfium outputs BGRA; convert to RGBA.
		img := image.NewRGBA(image.Rect(0, 0, pxW, pxH))
		for y := 0; y < pxH; y++ {
			for x := 0; x < pxW; x++ {
				off := y*stride + x*4
				img.SetRGBA(x, y, color.RGBA{
					R: pixels[off+2], // B
					G: pixels[off+1], // G
					B: pixels[off],   // R
					A: 255,
				})
			}
		}
		out = img
	})
	return
}

// openPage opens a document and page, returning post-rotation dimensions
// and a cleanup function.  Callers must call closeAll() to free resources.
func openPage(pdfData []byte, pageIdx int) (
	doc C.FPDF_DOCUMENT,
	page C.FPDF_PAGE,
	pw, ph float64,
	closeAll func(),
	err error,
) {
	cData := C.CBytes(pdfData)

	doc = C.FPDF_LoadMemDocument(unsafe.Pointer(cData), C.int(len(pdfData)), nil)
	if doc == nil {
		C.free(cData)
		err = fmt.Errorf("pdfium: FPDF_LoadMemDocument returned nil")
		return
	}

	page = C.FPDF_LoadPage(doc, C.int(pageIdx))
	if page == nil {
		C.FPDF_CloseDocument(doc)
		C.free(cData)
		err = fmt.Errorf("pdfium: FPDF_LoadPage(%d) returned nil", pageIdx)
		return
	}

	pw = float64(C.FPDF_GetPageWidth(page))
	ph = float64(C.FPDF_GetPageHeight(page))
	if pw <= 0 || ph <= 0 {
		C.FPDF_ClosePage(page)
		C.FPDF_CloseDocument(doc)
		C.free(cData)
		err = fmt.Errorf("pdfium: invalid page dimensions %.1fx%.1f", pw, ph)
		return
	}

	closeAll = func() {
		C.FPDF_ClosePage(page)
		C.FPDF_CloseDocument(doc)
		C.free(cData)
	}
	return
}

// ExtractOutlines returns the document outline (bookmarks / table of contents)
// from a PDF. Returns nil for empty or broken PDFs. Title decoding uses UTF-16LE
// as required by pdfium's FPDFBookmark_GetTitle.
//
// Traversal is iterative with an explicit stack to avoid stack overflow on deep
// outline trees. pdfium is not thread-safe; all FPDF_* calls are serialized by
// the shared pdfsync.Mu.
func ExtractOutlines(pdfData []byte) []Outline {
	if len(pdfData) == 0 {
		return nil
	}
	var out []Outline
	pdfsync.With(func() {
		Init()

		cData := C.CBytes(pdfData)
		defer C.free(cData)

		doc := C.FPDF_LoadMemDocument(unsafe.Pointer(cData), C.int(len(pdfData)), nil)
		if doc == nil {
			return
		}
		defer C.FPDF_CloseDocument(doc)

		type frame struct {
			bm    C.FPDF_BOOKMARK
			level int
		}

		var result []Outline
		stack := []frame{{bm: C.FPDFBookmark_GetFirstChild(doc, nil), level: 0}}

		for len(stack) > 0 {
			top := stack[len(stack)-1]
			stack = stack[:len(stack)-1]
			if top.bm == nil {
				continue
			}

			// Title (UTF-16LE).
			title := bookmarkTitle(top.bm)

			// Page number.
			pageNum := 1 // default to page 1 if dest is unavailable
			if dest := C.FPDFBookmark_GetDest(doc, top.bm); dest != nil {
				pn := C.FPDFDest_GetDestPageIndex(doc, dest)
				if pn >= 0 {
					pageNum = int(pn) + 1 // pdfium returns 0-based
				}
			}

			result = append(result, Outline{Title: title, Level: top.level, PageNumber: pageNum})

			// Push siblings after children so children are processed first (pre-order).
			if sibling := C.FPDFBookmark_GetNextSibling(doc, top.bm); sibling != nil {
				stack = append(stack, frame{bm: sibling, level: top.level})
			}
			if child := C.FPDFBookmark_GetFirstChild(doc, top.bm); child != nil {
				stack = append(stack, frame{bm: child, level: top.level + 1})
			}
		}
		out = result
	})
	return out
}

// bookmarkTitle reads a bookmark's title (UTF-16LE) and converts to Go string.
func bookmarkTitle(bm C.FPDF_BOOKMARK) string {
	// First call: get required buffer length in bytes.
	buflen := C.FPDFBookmark_GetTitle(bm, nil, 0)
	if buflen <= 0 {
		return ""
	}
	buf := make([]byte, buflen)
	C.FPDFBookmark_GetTitle(bm, unsafe.Pointer(&buf[0]), buflen)

	// Title is UTF-16LE. Convert to []uint16 then decode.
	n := int(buflen) / 2
	u16 := unsafe.Slice((*uint16)(unsafe.Pointer(&buf[0])), n)

	// Strip trailing null terminator if present.
	if n > 0 && u16[n-1] == 0 {
		u16 = u16[:n-1]
	}
	return string(utf16.Decode(u16))
}