petal/internal/docs/search.go

package docs

import (
	"net/http"
	"strings"
	"unicode/utf8"

	"github.com/go-chi/chi/v5"

	"gitea.parodia.dev/drwily/petal/internal/auth"
	"gitea.parodia.dev/drwily/petal/internal/db"
	"gitea.parodia.dev/drwily/petal/internal/httputil"
)

// Search snippet shaping.
const (
	// ftsMinRunes is the shortest query the trigram FTS index can match. Shorter
	// queries (common for 2-character Chinese words) fall back to a LIKE scan.
	ftsMinRunes = 3

	// snippetContext is how many runes of context to show on each side of the
	// matched term in a result snippet.
	snippetContext = 28

	// maxSearchResults caps how many hits we return — plenty for a personal
	// corpus, and keeps the response small.
	maxSearchResults = 50

	// hlStart/hlEnd wrap the matched span in a snippet. They're control-character
	// sentinels that never occur in real text, so the client can split on them to
	// highlight the match without escaping user content.
	hlStart = "\x01"
	hlEnd   = "\x02"
)

// searchResult is one hit: a document summary plus a highlighted snippet showing
// where the query matched.
type searchResult struct {
	ID        string   `json:"id"`
	Title     string   `json:"title"`
	WordCount int      `json:"word_count"`
	UpdatedAt string   `json:"updated_at"`
	Snippet   string   `json:"snippet"`
	Tags      []db.Tag `json:"tags"`
}

// SearchRoutes returns the router mounted at /api/search.
func (h *Handler) SearchRoutes() chi.Router {
	r := chi.NewRouter()
	r.Get("/", h.search)
	return r
}

// search runs a cross-document full-text search for the caller. Queries of
// three or more runes use the trigram FTS index (fast, ranked); shorter queries
// fall back to a LIKE scan so 2-character Chinese words still resolve. Either way
// the snippet is built in Go from the original text, for clean word boundaries
// and a uniform highlight format.
func (h *Handler) search(w http.ResponseWriter, r *http.Request) {
	userID := auth.UserID(r.Context())
	q := strings.TrimSpace(r.URL.Query().Get("q"))
	if q == "" {
		httputil.WriteJSON(w, http.StatusOK, []searchResult{})
		return
	}

	type row struct {
		id, title, contentText, updatedAt string
		wordCount                         int
	}
	var rows []row

	if utf8.RuneCountInString(q) >= ftsMinRunes {
		// Wrap the whole query as one FTS phrase (doubling embedded quotes), so
		// special characters are treated literally and trigram does a contiguous
		// substring match.
		phrase := `"` + strings.ReplaceAll(q, `"`, `""`) + `"`
		sqlRows, err := h.DB.Query(
			`SELECT d.id, d.title, d.content_text, d.word_count, d.updated_at
			   FROM documents_fts f
			   JOIN documents d ON d.id = f.doc_id
			  WHERE documents_fts MATCH ? AND d.user_id = ?
			  ORDER BY rank
			  LIMIT ?`,
			phrase, userID, maxSearchResults,
		)
		if err != nil {
			httputil.ServerError(w, err)
			return
		}
		defer sqlRows.Close()
		for sqlRows.Next() {
			var rw row
			if err := sqlRows.Scan(&rw.id, &rw.title, &rw.contentText, &rw.wordCount, &rw.updatedAt); err != nil {
				httputil.ServerError(w, err)
				return
			}
			rows = append(rows, rw)
		}
		if err := sqlRows.Err(); err != nil {
			httputil.ServerError(w, err)
			return
		}
	} else {
		// Short query: LIKE scan over title + body. Escape LIKE wildcards so a
		// literal % or _ in the query matches itself.
		like := "%" + escapeLike(q) + "%"
		sqlRows, err := h.DB.Query(
			`SELECT id, title, content_text, word_count, updated_at
			   FROM documents
			  WHERE user_id = ?
			    AND (title LIKE ? ESCAPE '\' OR content_text LIKE ? ESCAPE '\')
			  ORDER BY updated_at DESC
			  LIMIT ?`,
			userID, like, like, maxSearchResults,
		)
		if err != nil {
			httputil.ServerError(w, err)
			return
		}
		defer sqlRows.Close()
		for sqlRows.Next() {
			var rw row
			if err := sqlRows.Scan(&rw.id, &rw.title, &rw.contentText, &rw.wordCount, &rw.updatedAt); err != nil {
				httputil.ServerError(w, err)
				return
			}
			rows = append(rows, rw)
		}
		if err := sqlRows.Err(); err != nil {
			httputil.ServerError(w, err)
			return
		}
	}

	out := make([]searchResult, 0, len(rows))
	ids := make([]string, 0, len(rows))
	for _, rw := range rows {
		out = append(out, searchResult{
			ID:        rw.id,
			Title:     rw.title,
			WordCount: rw.wordCount,
			UpdatedAt: rw.updatedAt,
			Snippet:   buildSnippet(rw.title, rw.contentText, q),
		})
		ids = append(ids, rw.id)
	}

	byDoc, err := h.tagsByDoc(userID, ids)
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	for i := range out {
		out[i].Tags = byDoc[out[i].ID]
		if out[i].Tags == nil {
			out[i].Tags = []db.Tag{}
		}
	}
	httputil.WriteJSON(w, http.StatusOK, out)
}

// escapeLike escapes the LIKE metacharacters (% and _) and the escape character
// itself so the query is matched literally. Pairs with `ESCAPE '\'`.
func escapeLike(s string) string {
	r := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
	return r.Replace(s)
}

// buildSnippet returns a short excerpt around the first case-insensitive match of
// query, with the matched span wrapped in the hl sentinels. It prefers a body
// match (with surrounding context); if the query only appears in the title it
// highlights the title instead; otherwise it shows the body's opening so a result
// always shows something. Windowing is rune-aware so CJK is never split
// mid-character.
func buildSnippet(title, body, query string) string {
	runes := []rune(body)
	qLen := utf8.RuneCountInString(query)
	if idx := runeIndexFold(runes, query); idx >= 0 {
		start := idx - snippetContext
		if start < 0 {
			start = 0
		}
		end := idx + qLen + snippetContext
		if end > len(runes) {
			end = len(runes)
		}
		var b strings.Builder
		if start > 0 {
			b.WriteString("…")
		}
		b.WriteString(string(runes[start:idx]))
		b.WriteString(hlStart)
		b.WriteString(string(runes[idx : idx+qLen]))
		b.WriteString(hlEnd)
		b.WriteString(string(runes[idx+qLen : end]))
		if end < len(runes) {
			b.WriteString("…")
		}
		return b.String()
	}

	// Body had no literal match (title-only hit, or an FTS span the literal scan
	// can't reproduce). Highlight the title if the query is there.
	tRunes := []rune(title)
	if idx := runeIndexFold(tRunes, query); idx >= 0 {
		return string(tRunes[:idx]) + hlStart + string(tRunes[idx:idx+qLen]) + hlEnd + string(tRunes[idx+qLen:])
	}

	// Last resort: the body's opening as plain context.
	return clip(runes, 0, 2*snippetContext+qLen)
}

// clip returns runes[start:start+n] (bounded), with a trailing ellipsis when the
// body continues. Used for the no-direct-match fallback.
func clip(runes []rune, start, n int) string {
	if start >= len(runes) {
		return ""
	}
	end := start + n
	trailing := ""
	if end < len(runes) {
		trailing = "…"
	} else {
		end = len(runes)
	}
	return string(runes[start:end]) + trailing
}

// runeIndexFold finds the first index (in runes) where query occurs in runes,
// case-insensitively. Returns -1 if absent. Simple O(n*m) scan — fine for
// single-document snippet building.
func runeIndexFold(runes []rune, query string) int {
	q := []rune(strings.ToLower(query))
	if len(q) == 0 {
		return -1
	}
	lower := make([]rune, len(runes))
	for i, r := range runes {
		lower[i] = toLowerRune(r)
	}
	for i := 0; i+len(q) <= len(lower); i++ {
		match := true
		for j := 0; j < len(q); j++ {
			if lower[i+j] != q[j] {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}

// toLowerRune lowercases ASCII letters (the only case-bearing script here); CJK
// and other runes pass through unchanged.
func toLowerRune(r rune) rune {
	if r >= 'A' && r <= 'Z' {
		return r + ('a' - 'A')
	}
	return r
}