llmgw/internal/kb/extract/rtf.go

package extract

import (
	"errors"
	"strconv"
	"strings"
	"unicode/utf16"

	"golang.org/x/text/encoding/charmap"
	"golang.org/x/text/unicode/norm"
)

// rtfSkip are destinations whose content is not the document's text.
var rtfSkip = map[string]bool{
	"fonttbl": true, "colortbl": true, "stylesheet": true, "info": true, "pict": true,
	"header": true, "headerl": true, "headerr": true, "headerf": true, "footer": true,
	"footerl": true, "footerr": true, "footerf": true, "object": true, "xmlnstbl": true,
	"listtable": true, "listoverridetable": true, "rsidtbl": true, "generator": true,
	"themedata": true, "colorschememapping": true, "latentstyles": true, "datastore": true,
	"filetbl": true, "revtbl": true, "pgdsctbl": true, "fldinst": true, "bkmkstart": true,
	"bkmkend": true, "shppict": true, "nonshppict": true, "datafield": true, "blipuid": true,
}

// rtfText reads Rich Text Format: text, paragraphs, tabs and table cells,
// characters in the document's code page (\ansicpg) or as \u escapes;
// font tables, pictures, headers and other destinations skipped.
func rtfText(b []byte) (string, error) {
	s := string(b)
	if !strings.HasPrefix(strings.TrimSpace(s), "{\\rtf") {
		return "", errors.New("not RTF: no {\\rtf header")
	}
	type state struct {
		skip bool
		uc   int // how many fallback characters follow a \u
	}
	cp := charmap.Windows1252
	stack := []state{{uc: 1}}
	var o out
	var pending []byte // \'hh bytes, decoded in runs so a code page's pairs stay together
	var surrogate rune
	flush := func() {
		if len(pending) == 0 {
			return
		}
		d, _ := cp.NewDecoder().Bytes(pending)
		o.raw(string(d))
		pending = pending[:0]
	}
	skipChars := 0
	cur := func() *state { return &stack[len(stack)-1] }
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch c {
		case '{':
			flush()
			stack = append(stack, *cur())
		case '}':
			flush()
			if len(stack) > 1 {
				stack = stack[:len(stack)-1]
			}
		case '\\':
			if i+1 >= len(s) {
				break
			}
			n := s[i+1]
			switch {
			case n == '\'' && i+3 < len(s):
				v, err := strconv.ParseUint(s[i+2:i+4], 16, 8)
				i += 3
				if err != nil || cur().skip {
					continue
				}
				if skipChars > 0 {
					skipChars--
					continue
				}
				pending = append(pending, byte(v))
			case n == '*':
				i++
				cur().skip = true
			case n == '\\' || n == '{' || n == '}':
				i++
				if !cur().skip {
					flush()
					o.raw(string(n))
				}
			case n == '~':
				i++
				if !cur().skip {
					o.raw(" ")
				}
			case n == '-' || n == '_':
				i++
			case n == '\n' || n == '\r':
				i++
				if !cur().skip {
					flush()
					o.para()
				}
			case isLetter(n):
				j := i + 1
				for j < len(s) && isLetter(s[j]) {
					j++
				}
				word := s[i+1 : j]
				k := j
				if k < len(s) && (s[k] == '-' || (s[k] >= '0' && s[k] <= '9')) {
					k++
					for k < len(s) && s[k] >= '0' && s[k] <= '9' {
						k++
					}
				}
				arg, hasArg := 0, k > j
				if hasArg {
					arg, _ = strconv.Atoi(s[j:k])
				}
				if k < len(s) && s[k] == ' ' {
					k++
				}
				i = k - 1
				flush()
				if rtfSkip[word] {
					cur().skip = true
					continue
				}
				switch word {
				case "ansicpg":
					cp = codePage(arg)
				case "uc":
					cur().uc = arg
				case "u":
					if cur().skip {
						continue
					}
					r := rune(arg)
					if r < 0 {
						r += 65536
					}
					switch {
					case utf16.IsSurrogate(r) && r < 0xDC00:
						surrogate = r
					case utf16.IsSurrogate(r):
						o.raw(string(utf16.DecodeRune(surrogate, r)))
						surrogate = 0
					default:
						o.raw(string(r))
					}
					skipChars = cur().uc
				case "par", "sect", "page":
					if !cur().skip {
						o.para()
					}
				case "line":
					if !cur().skip {
						o.newline()
					}
				case "tab":
					if !cur().skip {
						o.raw("\t")
					}
				case "cell":
					if !cur().skip {
						o.tab()
					}
				case "row":
					if !cur().skip {
						o.newline()
					}
				case "emdash":
					o.raw("—")
				case "endash":
					o.raw("–")
				case "bullet":
					o.raw("•")
				case "lquote", "rquote":
					o.raw("'")
				case "ldblquote", "rdblquote":
					o.raw("\"")
				}
			default:
				i++
			}
		case '\r', '\n':
			// Line breaks in RTF source are not text.
		default:
			if cur().skip {
				continue
			}
			if skipChars > 0 {
				skipChars--
				continue
			}
			if c >= 0x80 {
				// Unescaped 8-bit text, as some writers leave it.
				pending = append(pending, c)
				continue
			}
			flush()
			o.raw(string(c))
		}
	}
	flush()
	return norm.NFC.String(o.String()), nil
}

func isLetter(c byte) bool { return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') }