tm_back/pkg/mongo/cursor.go

package mongo

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

	"go.mongodb.org/mongo-driver/v2/bson"
)

const pageCursorVersion = 1

// pageCursor is an opaque pagination token for keyset (cursor) pagination.
type pageCursor struct {
	V           int    `json:"v"`
	SortField   string `json:"sort_field"`
	SortOrder   int    `json:"sort_order"`
	ValueType   string `json:"value_type"` // i64, f64, str, oid
	SortValue   string `json:"sort_value"`
	ID          string `json:"id"`
	FiltersHash string `json:"filters_hash,omitempty"`
	PageOffset  int    `json:"page_offset,omitempty"`
}

func encodePageCursor(sortField string, sortOrder int, sortValue interface{}, id, filtersHash string, pageOffset int) (string, error) {
	if sortField == "" {
		return "", fmt.Errorf("%w: sort field is required", ErrInvalidCursor)
	}
	if id == "" {
		return "", fmt.Errorf("%w: document id is required", ErrInvalidCursor)
	}

	valueType, sortValueStr, err := formatCursorSortValue(sortField, sortValue)
	if err != nil {
		return "", err
	}

	payload, err := json.Marshal(pageCursor{
		V:           pageCursorVersion,
		SortField:   sortField,
		SortOrder:   sortOrder,
		ValueType:   valueType,
		SortValue:   sortValueStr,
		ID:          id,
		FiltersHash: filtersHash,
		PageOffset:  pageOffset,
	})
	if err != nil {
		return "", fmt.Errorf("failed to marshal cursor: %w", err)
	}

	return base64.RawURLEncoding.EncodeToString(payload), nil
}

func decodePageCursor(encoded string) (*pageCursor, error) {
	cursorBytes, err := base64.RawURLEncoding.DecodeString(encoded)
	if err != nil {
		// Accept standard Base64 cursors issued before URL-safe encoding.
		cursorBytes, err = base64.StdEncoding.DecodeString(encoded)
	}
	if err != nil {
		return nil, fmt.Errorf("%w: invalid cursor encoding", ErrInvalidCursor)
	}

	var cursor pageCursor
	if err := json.Unmarshal(cursorBytes, &cursor); err != nil {
		return nil, fmt.Errorf("%w: invalid cursor payload", ErrInvalidCursor)
	}
	if cursor.V != pageCursorVersion || cursor.SortField == "" || cursor.ID == "" {
		return nil, fmt.Errorf("%w: unsupported or incomplete cursor", ErrInvalidCursor)
	}
	if cursor.SortOrder != 1 && cursor.SortOrder != -1 {
		return nil, fmt.Errorf("%w: invalid sort order in cursor", ErrInvalidCursor)
	}

	return &cursor, nil
}

func (c *pageCursor) matches(sortField string, sortOrder int) bool {
	return c.SortField == sortField && c.SortOrder == sortOrder
}

func (c *pageCursor) mongoFilter() (bson.M, error) {
	if c.SortField == "_id" {
		objectID, err := bson.ObjectIDFromHex(c.ID)
		if err != nil {
			return nil, fmt.Errorf("%w: invalid ObjectID in cursor", ErrInvalidCursor)
		}
		operator := "$lt"
		if c.SortOrder == 1 {
			operator = "$gt"
		}
		return bson.M{"_id": bson.M{operator: objectID}}, nil
	}

	sortValue, err := c.parseSortValue()
	if err != nil {
		return nil, err
	}

	objectID, err := bson.ObjectIDFromHex(c.ID)
	if err != nil {
		return nil, fmt.Errorf("%w: invalid ObjectID in cursor", ErrInvalidCursor)
	}

	primaryOp := "$lt"
	secondaryOp := "$lt"
	if c.SortOrder == 1 {
		primaryOp = "$gt"
		secondaryOp = "$gt"
	}

	return bson.M{
		"$or": []bson.M{
			{c.SortField: bson.M{primaryOp: sortValue}},
			{
				c.SortField: sortValue,
				"_id":       bson.M{secondaryOp: objectID},
			},
		},
	}, nil
}

func (c *pageCursor) parseSortValue() (interface{}, error) {
	switch c.ValueType {
	case "i64":
		v, err := strconv.ParseInt(c.SortValue, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("%w: invalid int64 sort value", ErrInvalidCursor)
		}
		return v, nil
	case "f64":
		v, err := strconv.ParseFloat(c.SortValue, 64)
		if err != nil {
			return nil, fmt.Errorf("%w: invalid float sort value", ErrInvalidCursor)
		}
		return v, nil
	case "str":
		return c.SortValue, nil
	case "oid":
		return bson.ObjectIDFromHex(c.SortValue)
	default:
		return nil, fmt.Errorf("%w: unknown sort value type", ErrInvalidCursor)
	}
}

func formatCursorSortValue(sortField string, sortValue interface{}) (string, string, error) {
	if sortField == "_id" {
		switch v := sortValue.(type) {
		case bson.ObjectID:
			return "oid", v.Hex(), nil
		case string:
			if _, err := bson.ObjectIDFromHex(v); err != nil {
				return "", "", fmt.Errorf("%w: invalid ObjectID sort value", ErrInvalidCursor)
			}
			return "oid", v, nil
		default:
			return "", "", fmt.Errorf("%w: unsupported _id sort value", ErrInvalidCursor)
		}
	}

	switch v := sortValue.(type) {
	case int64:
		return "i64", strconv.FormatInt(v, 10), nil
	case int32:
		return "i64", strconv.FormatInt(int64(v), 10), nil
	case int:
		return "i64", strconv.FormatInt(int64(v), 10), nil
	case float64:
		return "f64", strconv.FormatFloat(v, 'f', -1, 64), nil
	case float32:
		return "f64", strconv.FormatFloat(float64(v), 'f', -1, 64), nil
	case string:
		return "str", v, nil
	case bson.ObjectID:
		return "oid", v.Hex(), nil
	default:
		return "", "", fmt.Errorf("%w: unsupported sort value type for field %s", ErrInvalidCursor, sortField)
	}
}

func documentSortValue(docMap bson.M, sortField string) (interface{}, error) {
	if sortField == "_id" {
		id, exists := docMap["_id"]
		if !exists {
			return nil, fmt.Errorf("sort field _id not found in document")
		}
		return id, nil
	}

	if strings.Contains(sortField, ".") {
		return nestedDocumentValue(docMap, sortField)
	}

	sortValue, exists := docMap[sortField]
	if !exists {
		return nil, fmt.Errorf("sort field %s not found in document", sortField)
	}
	return sortValue, nil
}

func nestedDocumentValue(docMap bson.M, path string) (interface{}, error) {
	parts := strings.Split(path, ".")
	var current interface{} = docMap

	for _, part := range parts {
		nested, ok := current.(bson.M)
		if !ok {
			if nestedMap, ok := current.(map[string]interface{}); ok {
				nested = bson.M(nestedMap)
			} else {
				return nil, fmt.Errorf("sort field %s not found in document", path)
			}
		}

		value, exists := nested[part]
		if !exists {
			return nil, fmt.Errorf("sort field %s not found in document", path)
		}
		current = value
	}

	return current, nil
}

func documentIDHex(docMap bson.M) (string, error) {
	id, exists := docMap["_id"]
	if !exists {
		return "", fmt.Errorf("document _id not found")
	}
	switch v := id.(type) {
	case bson.ObjectID:
		return v.Hex(), nil
	case string:
		if _, err := bson.ObjectIDFromHex(v); err != nil {
			return "", fmt.Errorf("invalid document _id")
		}
		return v, nil
	default:
		return "", fmt.Errorf("unsupported document _id type")
	}
}

func cloneFilter(filter bson.M) (bson.M, error) {
	if len(filter) == 0 {
		return bson.M{}, nil
	}
	raw, err := bson.Marshal(filter)
	if err != nil {
		return nil, err
	}
	var cloned bson.M
	if err := bson.Unmarshal(raw, &cloned); err != nil {
		return nil, err
	}
	return cloned, nil
}

func mergeFilters(base, extra bson.M) bson.M {
	if len(extra) == 0 {
		return base
	}
	if len(base) == 0 {
		return extra
	}
	return bson.M{
		"$and": []bson.M{base, extra},
	}
}