Warp/server/internal/syncapi/search.go

package syncapi

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/dariuszjaniak/warp/server/internal/roomserver"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// POST /search: full-text search of the rooms a user is or was in.
//
// The roomserver decides where the user may look and the store runs the
// query; what is here is the request and response shape, which carries a
// fair amount besides the hits — the events around each one, the profiles of
// their senders, the rooms' state, and the hits grouped by room or sender —
// all of it optional and all of it what a client needs to show a result
// without a second round of requests.

// DefaultSearchLimit is the page size when the filter names none.
const DefaultSearchLimit = 10

// MaxSearchLimit caps a page. Every hit can bring its context, its
// senders' profiles and its room's state with it, so a page is far heavier
// than its count suggests.
const MaxSearchLimit = 100

// MaxSearchContext caps before_limit and after_limit.
const MaxSearchContext = 20

// SearchRequest is the body of a /search, with the caller.
type SearchRequest struct {
	User string

	// NextBatch is the ?next_batch= of a later page.
	NextBatch string

	Body SearchBody
}

// SearchBody is the request body.
type SearchBody struct {
	SearchCategories struct {
		RoomEvents *RoomEventsCriteria `json:"room_events"`
	} `json:"search_categories"`
}

// RoomEventsCriteria is the one search category the spec defines.
type RoomEventsCriteria struct {
	SearchTerm   string           `json:"search_term"`
	Keys         []string         `json:"keys,omitempty"`
	Filter       *RoomEventFilter `json:"filter,omitempty"`
	OrderBy      string           `json:"order_by,omitempty"`
	EventContext *struct {
		BeforeLimit    *int `json:"before_limit,omitempty"`
		AfterLimit     *int `json:"after_limit,omitempty"`
		IncludeProfile bool `json:"include_profile,omitempty"`
	} `json:"event_context,omitempty"`
	IncludeState bool `json:"include_state,omitempty"`
	Groupings    *struct {
		GroupBy []struct {
			Key string `json:"key"`
		} `json:"group_by"`
	} `json:"groupings,omitempty"`
}

// SearchResponse is the body of the answer.
type SearchResponse struct {
	SearchCategories struct {
		RoomEvents RoomEventsResults `json:"room_events"`
	} `json:"search_categories"`
}

// RoomEventsResults is one page of hits.
type RoomEventsResults struct {
	Count      int                                    `json:"count"`
	Highlights []string                               `json:"highlights"`
	Results    []SearchHit                            `json:"results"`
	NextBatch  string                                 `json:"next_batch,omitempty"`
	State      map[string][]roomserver.ClientEvent    `json:"state,omitempty"`
	Groups     map[string]map[string]SearchGroupEntry `json:"groups,omitempty"`
}

// SearchHit is one result.
type SearchHit struct {
	Rank    float64                `json:"rank"`
	Result  roomserver.ClientEvent `json:"result"`
	Context *SearchContext         `json:"context,omitempty"`
}

// SearchContext is what surrounds one hit.
type SearchContext struct {
	Start        string                       `json:"start,omitempty"`
	End          string                       `json:"end,omitempty"`
	EventsBefore []roomserver.ClientEvent     `json:"events_before"`
	EventsAfter  []roomserver.ClientEvent     `json:"events_after"`
	ProfileInfo  map[string]SearchProfileInfo `json:"profile_info,omitempty"`
}

// SearchProfileInfo is a sender's name and avatar.
type SearchProfileInfo struct {
	DisplayName string `json:"displayname,omitempty"`
	AvatarURL   string `json:"avatar_url,omitempty"`
}

// SearchGroupEntry is one group of a grouped search.
type SearchGroupEntry struct {
	Results   []string `json:"results"`
	Order     int      `json:"order"`
	NextBatch string   `json:"next_batch,omitempty"`
}

// The search tokens. A ranked search pages by offset, because the ranking
// of a result does not move as later events arrive; a search by recency
// pages by stream position, which keeps a page from repeating a result when
// new messages land between two requests.
const (
	rankTokenPrefix   = "r"
	recentTokenPrefix = "p"
)

// Search answers one /search.
func (s *API) Search(ctx context.Context, req SearchRequest) (SearchResponse, error) {
	criteria := req.Body.SearchCategories.RoomEvents
	var resp SearchResponse
	resp.SearchCategories.RoomEvents.Highlights = []string{}
	resp.SearchCategories.RoomEvents.Results = []SearchHit{}
	if criteria == nil {
		// The only category there is was not asked for, so there is
		// nothing to search.
		return resp, nil
	}
	if strings.TrimSpace(criteria.SearchTerm) == "" {
		return SearchResponse{}, errInvalidParamf("search_term must not be empty")
	}

	order := storage.SearchByRank
	switch criteria.OrderBy {
	case "", "rank":
	case "recent":
		order = storage.SearchByRecent
	default:
		return SearchResponse{}, errInvalidParamf("order_by must be \"rank\" or \"recent\"")
	}
	for _, key := range criteria.Keys {
		switch key {
		case "content.body", "content.name", "content.topic":
		default:
			return SearchResponse{}, errInvalidParamf("%q is not a search key", key)
		}
	}

	offset, before, err := parseSearchToken(req.NextBatch, order)
	if err != nil {
		return SearchResponse{}, err
	}

	limit := DefaultSearchLimit
	var filter storage.TimelineFilter
	var rooms, notRooms []string
	if f := criteria.Filter; f != nil {
		if f.Limit != nil && *f.Limit > 0 {
			limit = min(*f.Limit, MaxSearchLimit)
		}
		filter = roomEventFilter(f)
		rooms, notRooms = f.Rooms, f.NotRooms
	}

	found, total, err := s.rooms.Search(ctx, roomserver.SearchRequest{
		User:     req.User,
		Term:     criteria.SearchTerm,
		Keys:     criteria.Keys,
		Rooms:    rooms,
		NotRooms: notRooms,
		Filter:   filter,
		Order:    order,
		Offset:   offset,
		Before:   before,
		Limit:    limit,
	})
	if err != nil {
		return SearchResponse{}, err
	}

	out := &resp.SearchCategories.RoomEvents
	out.Count = total
	out.Highlights = storage.SearchWords(criteria.SearchTerm)

	for _, hit := range found {
		h := SearchHit{Rank: hit.Rank, Result: hit.Event}
		if criteria.EventContext != nil {
			if h.Context, err = s.searchContext(ctx, req.User, hit, criteria); err != nil {
				return SearchResponse{}, err
			}
		}
		out.Results = append(out.Results, h)
	}

	// There is more when the page came back full and the total says so.
	if len(found) == limit {
		switch order {
		case storage.SearchByRank:
			if offset+len(found) < total {
				out.NextBatch = rankTokenPrefix + strconv.Itoa(offset+len(found))
			}
		case storage.SearchByRecent:
			out.NextBatch = recentTokenPrefix + strconv.FormatInt(found[len(found)-1].Position, 10)
		}
	}

	if criteria.IncludeState {
		out.State = map[string][]roomserver.ClientEvent{}
		for _, hit := range found {
			roomID := hit.Event.RoomID
			if _, done := out.State[roomID]; done {
				continue
			}
			state, err := s.rooms.RoomState(ctx, roomID, req.User)
			if err != nil {
				return SearchResponse{}, err
			}
			out.State[roomID] = state
		}
	}

	if criteria.Groupings != nil {
		out.Groups = searchGroups(found, criteria, out.NextBatch)
	}
	return resp, nil
}

// searchContext reads the events around one hit.
func (s *API) searchContext(ctx context.Context, userID string, hit roomserver.SearchResult, criteria *RoomEventsCriteria) (*SearchContext, error) {
	limitOf := func(p *int) int {
		if p == nil {
			return 5
		}
		return max(0, min(*p, MaxSearchContext))
	}
	beforeLimit := limitOf(criteria.EventContext.BeforeLimit)
	afterLimit := limitOf(criteria.EventContext.AfterLimit)
	roomID := hit.Event.RoomID

	out := &SearchContext{
		EventsBefore: []roomserver.ClientEvent{},
		EventsAfter:  []roomserver.ClientEvent{},
	}
	if beforeLimit > 0 {
		earlier, err := s.rooms.Timeline(ctx, roomserver.TimelineRequest{
			RoomID:    roomID,
			User:      userID,
			High:      storage.PositionBefore(hit.Position),
			Direction: storage.DirectionBackward,
			Limit:     beforeLimit,
		})
		if err != nil {
			return nil, err
		}
		if len(earlier.Events) > 0 {
			out.EventsBefore = reversed(earlier.Events)
			out.Start = pduToken(earlier.Start).String()
		}
	}
	if afterLimit > 0 {
		later, err := s.rooms.Timeline(ctx, roomserver.TimelineRequest{
			RoomID:    roomID,
			User:      userID,
			Low:       hit.Position,
			Direction: storage.DirectionForward,
			Limit:     afterLimit,
		})
		if err != nil {
			return nil, err
		}
		if len(later.Events) > 0 {
			out.EventsAfter = later.Events
			out.End = pduToken(later.End).String()
		}
	}

	if criteria.EventContext.IncludeProfile {
		senders := sendersOf(append(append([]roomserver.ClientEvent{hit.Event}, out.EventsBefore...), out.EventsAfter...))
		members, err := s.rooms.MemberEvents(ctx, roomID, senders)
		if err != nil {
			return nil, err
		}
		out.ProfileInfo = map[string]SearchProfileInfo{}
		for _, m := range members {
			if m.StateKey == nil {
				continue
			}
			var content struct {
				DisplayName string `json:"displayname"`
				AvatarURL   string `json:"avatar_url"`
			}
			_ = json.Unmarshal(m.Content, &content)
			out.ProfileInfo[*m.StateKey] = SearchProfileInfo{DisplayName: content.DisplayName, AvatarURL: content.AvatarURL}
		}
	}
	return out, nil
}

// searchGroups groups a page of hits by room or sender, in the order the
// page ranked them.
func searchGroups(found []roomserver.SearchResult, criteria *RoomEventsCriteria, nextBatch string) map[string]map[string]SearchGroupEntry {
	out := map[string]map[string]SearchGroupEntry{}
	for _, g := range criteria.Groupings.GroupBy {
		var keyOf func(roomserver.SearchResult) string
		switch g.Key {
		case "room_id":
			keyOf = func(r roomserver.SearchResult) string { return r.Event.RoomID }
		case "sender":
			keyOf = func(r roomserver.SearchResult) string { return r.Event.Sender }
		default:
			continue
		}
		groups := map[string]SearchGroupEntry{}
		for _, hit := range found {
			key := keyOf(hit)
			entry, ok := groups[key]
			if !ok {
				entry.Order = len(groups) + 1
				entry.NextBatch = nextBatch
			}
			entry.Results = append(entry.Results, hit.Event.EventID)
			groups[key] = entry
		}
		out[g.Key] = groups
	}
	return out
}

// parseSearchToken reads a ?next_batch=.
func parseSearchToken(raw string, order storage.SearchOrder) (offset int, before int64, err error) {
	if raw == "" {
		return 0, 0, nil
	}
	bad := errInvalidTokenf("%q is not a search token", raw)
	switch order {
	case storage.SearchByRecent:
		body, ok := strings.CutPrefix(raw, recentTokenPrefix)
		if !ok {
			return 0, 0, bad
		}
		if before, err = strconv.ParseInt(body, 10, 64); err != nil {
			return 0, 0, bad
		}
		return 0, before, nil
	default:
		body, ok := strings.CutPrefix(raw, rankTokenPrefix)
		if !ok {
			return 0, 0, bad
		}
		if offset, err = strconv.Atoi(body); err != nil || offset < 0 {
			return 0, 0, fmt.Errorf("%w", bad)
		}
		return offset, 0, nil
	}
}