tm_back/pkg/mongo/repository.go

package mongo

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

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

// Common errors
var (
	ErrDocumentNotFound  = errors.New("document not found")
	ErrInvalidCursor     = errors.New("invalid cursor")
	ErrInvalidPagination = errors.New("invalid pagination parameters")
)

// Pagination represents pagination parameters for queries
type Pagination struct {
	Limit     int    `json:"limit"`     // Number of documents to return
	Skip      int    `json:"skip"`      // Number of documents to skip (for offset-based)
	Cursor    string `json:"cursor"`    // Base64 encoded cursor for cursor-based pagination
	SortField string `json:"sortField"` // Field to sort by (default: "_id")
	SortOrder int    `json:"sortOrder"` // Sort order: 1 for ascending, -1 for descending
	// SkipCount skips CountDocuments (use with cursor or deep offset pagination).
	SkipCount bool `json:"skipCount,omitempty"`
	// Projection limits fields returned from Find (e.g. exclude large blobs on list endpoints).
	Projection bson.M `json:"projection,omitempty"`
}

// PaginatedResult represents a paginated result set
type PaginatedResult[T any] struct {
	Items      []T    `json:"items"`      // The actual documents
	TotalCount int64  `json:"totalCount"` // Total number of documents matching the filter
	NextCursor string `json:"nextCursor"` // Base64 encoded cursor for next page
	HasMore    bool   `json:"hasMore"`    // Whether there are more documents
	PageOffset int    `json:"pageOffset"` // Zero-based offset of the first item in this page
}

// Repository defines the interface for MongoDB operations
type Repository[T any] interface {
	// FindByID retrieves a document by its ID
	FindByID(ctx context.Context, id string) (*T, error)

	// FindAll retrieves documents with optional filtering and pagination
	FindAll(ctx context.Context, filter bson.M, pagination Pagination) (*PaginatedResult[T], error)

	// Create inserts a new document
	Create(ctx context.Context, model *T) error

	// CreateMany inserts multiple documents
	CreateMany(ctx context.Context, models []T) error

	// Update updates an existing document
	Update(ctx context.Context, model *T) error

	// Delete removes a document by ID
	Delete(ctx context.Context, id string) error

	// Count returns the number of documents matching the filter
	Count(ctx context.Context, filter bson.M) (int64, error)

	// Aggregate executes an aggregation pipeline
	Aggregate(ctx context.Context, pipeline mongo.Pipeline) ([]bson.M, error)

	// FindOne retrieves a single document matching the filter
	FindOne(ctx context.Context, filter bson.M) (*T, error)

	// FindMany retrieves multiple documents without pagination
	FindMany(ctx context.Context, filter bson.M, limit int) ([]T, error)
}

// repository implements the Repository interface
type repository[T any] struct {
	collection *mongo.Collection
	logger     Logger
}

// NewRepository creates a new repository instance
func NewRepository[T any](collection *mongo.Collection, logger Logger) Repository[T] {
	return &repository[T]{
		collection: collection,
		logger:     logger,
	}
}

// FindByID retrieves a document by its ID
func (r *repository[T]) FindByID(ctx context.Context, id string) (*T, error) {
	objectID, err := bson.ObjectIDFromHex(id)
	if err != nil {
		r.logger.Error("Invalid ObjectID format", map[string]interface{}{
			"id":    id,
			"error": err.Error(),
		})
		return nil, fmt.Errorf("invalid id format: %w", err)
	}

	filter := bson.M{"_id": objectID}
	var result T

	err = r.collection.FindOne(ctx, filter).Decode(&result)
	if err != nil {
		if errors.Is(err, mongo.ErrNoDocuments) {
			return nil, ErrDocumentNotFound
		}
		r.logger.Error("Failed to find document by ID", map[string]interface{}{
			"id":    id,
			"error": err.Error(),
		})
		return nil, fmt.Errorf("failed to find document: %w", err)
	}

	return &result, nil
}

// FindAll retrieves documents with optional filtering and pagination
func (r *repository[T]) FindAll(ctx context.Context, filter bson.M, pagination Pagination) (*PaginatedResult[T], error) {
	// Validate pagination parameters
	if err := r.validatePagination(pagination); err != nil {
		return nil, err
	}

	// Set default values
	if pagination.Limit <= 0 {
		pagination.Limit = 10
	}
	if pagination.SortField == "" {
		pagination.SortField = "_id"
	}
	if pagination.SortOrder == 0 {
		pagination.SortOrder = -1 // Default to descending
	}

	// Build sort with _id tie-breaker for stable keyset pagination.
	sort := bson.D{{Key: pagination.SortField, Value: pagination.SortOrder}}
	if pagination.SortField != "_id" {
		sort = append(sort, bson.E{Key: "_id", Value: pagination.SortOrder})
	}

	countFilter, err := cloneFilter(filter)
	if err != nil {
		return nil, fmt.Errorf("failed to clone filter for count: %w", err)
	}

	filtersHash := HashFilter(filter)
	pageOffset := pagination.Skip
	if pagination.Cursor != "" {
		pageCursor, err := decodePageCursor(pagination.Cursor)
		if err != nil {
			return nil, err
		}
		pageOffset = pageCursor.PageOffset
	}

	// Build find options
	findOptions := options.Find().
		SetSort(sort).
		SetLimit(int64(pagination.Limit + 1)) // +1 to check if there are more documents

	if len(pagination.Projection) > 0 {
		findOptions.SetProjection(pagination.Projection)
	}

	// Handle cursor-based pagination
	if pagination.Cursor != "" {
		cursorFilter, err := r.buildCursorFilter(pagination.Cursor, pagination.SortField, pagination.SortOrder, filtersHash)
		if err != nil {
			return nil, err
		}
		filter = mergeFilters(filter, cursorFilter)
	} else if pagination.Skip > 0 {
		// Use offset-based pagination
		findOptions.SetSkip(int64(pagination.Skip))
	}

	// Run count in parallel with the find so list latency is max(find, count), not sum.
	var totalCount int64 = -1
	var countErr error
	var countWG sync.WaitGroup
	if !pagination.SkipCount {
		countWG.Add(1)
		go func() {
			defer countWG.Done()
			totalCount, countErr = r.Count(ctx, countFilter)
		}()
	}

	cursor, err := r.collection.Find(ctx, filter, findOptions)
	if err != nil {
		countWG.Wait()
		r.logger.Error("Failed to execute find query", map[string]interface{}{
			"filter":     filter,
			"pagination": pagination,
			"error":      err.Error(),
		})
		return nil, fmt.Errorf("failed to execute query: %w", err)
	}
	defer cursor.Close(ctx)

	var items []T
	if err = cursor.All(ctx, &items); err != nil {
		countWG.Wait()
		r.logger.Error("Failed to decode query results", map[string]interface{}{
			"error": err.Error(),
		})
		return nil, fmt.Errorf("failed to decode results: %w", err)
	}

	hasMore := false
	if len(items) > pagination.Limit {
		hasMore = true
		items = items[:pagination.Limit]
	}

	countWG.Wait()
	if !pagination.SkipCount {
		if countErr != nil {
			r.logger.Error("Failed to get total count", map[string]interface{}{
				"error": countErr.Error(),
			})
			totalCount = 0
		}
	}

	// Generate next cursor
	var nextCursor string
	if hasMore && len(items) > 0 {
		nextPageOffset := pageOffset + pagination.Limit
		nextCursor, err = r.generateCursor(items[len(items)-1], pagination.SortField, pagination.SortOrder, filtersHash, nextPageOffset)
		if err != nil {
			r.logger.Error("Failed to generate next cursor", map[string]interface{}{
				"error": err.Error(),
			})
			return nil, fmt.Errorf("failed to generate pagination cursor: %w", err)
		}
	}

	return &PaginatedResult[T]{
		Items:      items,
		TotalCount: totalCount,
		NextCursor: nextCursor,
		HasMore:    hasMore,
		PageOffset: pageOffset,
	}, nil
}

// Create inserts a new document
func (r *repository[T]) Create(ctx context.Context, model *T) error {
	// Set timestamps if the model implements Timestampable
	if timestampable, ok := any(model).(Timestampable); ok {
		now := time.Now().Unix()
		timestampable.SetCreatedAt(now)
		timestampable.SetUpdatedAt(now)
	}

	result, err := r.collection.InsertOne(ctx, model)
	if err != nil {
		r.logger.Error("Failed to create document", map[string]interface{}{
			"error": err.Error(),
		})
		return fmt.Errorf("failed to create document: %w", err)
	}

	// Set the ID if the model implements IDSetter
	if idSetter, ok := any(model).(IDSetter); ok {
		if oid, ok := result.InsertedID.(bson.ObjectID); ok {
			idSetter.SetID(oid.Hex())
		}
	}

	r.logger.Info("Document created successfully", map[string]interface{}{
		"id": result.InsertedID,
	})

	return nil
}

// CreateMany inserts multiple documents
func (r *repository[T]) CreateMany(ctx context.Context, models []T) error {
	// Set timestamps if the model implements Timestampable
	for _, v := range models {
		if timestampable, ok := any(v).(Timestampable); ok {
			now := time.Now().Unix()
			timestampable.SetCreatedAt(now)
			timestampable.SetUpdatedAt(now)
		}
	}

	result, err := r.collection.InsertMany(ctx, models)
	if err != nil {
		r.logger.Error("Failed to create multiple documents", map[string]interface{}{
			"error": err.Error(),
		})
		return fmt.Errorf("failed to create multiple documents: %w", err)
	}

	r.logger.Info("Documents created successfully", map[string]interface{}{
		"ids": result.InsertedIDs,
	})

	return nil
}

// Update updates an existing document
func (r *repository[T]) Update(ctx context.Context, model *T) error {
	// Set updated timestamp if the model implements Timestampable
	if timestampable, ok := any(model).(Timestampable); ok {
		timestampable.SetUpdatedAt(time.Now().Unix())
	}

	// Get the ID from the model
	id, err := r.getModelID(model)
	if err != nil {
		return err
	}

	objectID, err := bson.ObjectIDFromHex(id)
	if err != nil {
		return fmt.Errorf("invalid id format: %w", err)
	}

	filter := bson.M{"_id": objectID}
	update := bson.M{"$set": model}

	result, err := r.collection.UpdateOne(ctx, filter, update)
	if err != nil {
		r.logger.Error("Failed to update document", map[string]interface{}{
			"id":    id,
			"error": err.Error(),
		})
		return fmt.Errorf("failed to update document: %w", err)
	}

	if result.MatchedCount == 0 {
		return ErrDocumentNotFound
	}

	r.logger.Info("Document updated successfully", map[string]interface{}{
		"id": id,
	})

	return nil
}

// Delete removes a document by ID
func (r *repository[T]) Delete(ctx context.Context, id string) error {
	objectID, err := bson.ObjectIDFromHex(id)
	if err != nil {
		return fmt.Errorf("invalid id format: %w", err)
	}

	filter := bson.M{"_id": objectID}
	result, err := r.collection.DeleteOne(ctx, filter)
	if err != nil {
		r.logger.Error("Failed to delete document", map[string]interface{}{
			"id":    id,
			"error": err.Error(),
		})
		return fmt.Errorf("failed to delete document: %w", err)
	}

	if result.DeletedCount == 0 {
		return ErrDocumentNotFound
	}

	r.logger.Info("Document deleted successfully", map[string]interface{}{
		"id": id,
	})

	return nil
}

// isEmptyFilter reports whether a MongoDB filter matches all documents.
func isEmptyFilter(filter bson.M) bool {
	return len(filter) == 0
}

// Count returns the number of documents matching the filter.
// Unfiltered counts use EstimatedDocumentCount for O(1) metadata lookup on large collections.
func (r *repository[T]) Count(ctx context.Context, filter bson.M) (int64, error) {
	var count int64
	var err error
	if isEmptyFilter(filter) {
		count, err = r.collection.EstimatedDocumentCount(ctx)
	} else {
		count, err = r.collection.CountDocuments(ctx, filter)
	}
	if err != nil {
		r.logger.Error("Failed to count documents", map[string]interface{}{
			"filter": filter,
			"error":  err.Error(),
		})
		return 0, fmt.Errorf("failed to count documents: %w", err)
	}

	return count, nil
}

// Aggregate executes an aggregation pipeline
func (r *repository[T]) Aggregate(ctx context.Context, pipeline mongo.Pipeline) ([]bson.M, error) {
	cursor, err := r.collection.Aggregate(ctx, pipeline)
	if err != nil {
		r.logger.Error("Failed to execute aggregation", map[string]interface{}{
			"pipeline": pipeline,
			"error":    err.Error(),
		})
		return nil, fmt.Errorf("failed to execute aggregation: %w", err)
	}
	defer cursor.Close(ctx)

	var results []bson.M
	if err = cursor.All(ctx, &results); err != nil {
		r.logger.Error("Failed to decode aggregation results", map[string]interface{}{
			"error": err.Error(),
		})
		return nil, fmt.Errorf("failed to decode aggregation results: %w", err)
	}

	return results, nil
}

// FindOne retrieves a single document matching the filter
func (r *repository[T]) FindOne(ctx context.Context, filter bson.M) (*T, error) {
	var result T

	err := r.collection.FindOne(ctx, filter).Decode(&result)
	if err != nil {
		if errors.Is(err, mongo.ErrNoDocuments) {
			return nil, ErrDocumentNotFound
		}
		r.logger.Error("Failed to find document", map[string]interface{}{
			"filter": filter,
			"error":  err.Error(),
		})
		return nil, fmt.Errorf("failed to find document: %w", err)
	}

	return &result, nil
}

// FindMany retrieves multiple documents without pagination
func (r *repository[T]) FindMany(ctx context.Context, filter bson.M, limit int) ([]T, error) {
	findOptions := options.Find()
	if limit > 0 {
		findOptions.SetLimit(int64(limit))
	}

	cursor, err := r.collection.Find(ctx, filter, findOptions)
	if err != nil {
		r.logger.Error("Failed to execute find query", map[string]interface{}{
			"filter": filter,
			"limit":  limit,
			"error":  err.Error(),
		})
		return nil, fmt.Errorf("failed to execute query: %w", err)
	}
	defer cursor.Close(ctx)

	var results []T
	if err = cursor.All(ctx, &results); err != nil {
		r.logger.Error("Failed to decode query results", map[string]interface{}{
			"error": err.Error(),
		})
		return nil, fmt.Errorf("failed to decode results: %w", err)
	}

	return results, nil
}

// Helper methods

// validatePagination validates pagination parameters
func (r *repository[T]) validatePagination(pagination Pagination) error {
	if pagination.Limit < 0 {
		return fmt.Errorf("%w: limit cannot be negative", ErrInvalidPagination)
	}
	if pagination.Skip < 0 {
		return fmt.Errorf("%w: skip cannot be negative", ErrInvalidPagination)
	}
	if pagination.SortOrder != 0 && pagination.SortOrder != 1 && pagination.SortOrder != -1 {
		return fmt.Errorf("%w: sort order must be 1, -1, or 0", ErrInvalidPagination)
	}
	return nil
}

// buildCursorFilter builds a filter for cursor-based pagination.
func (r *repository[T]) buildCursorFilter(cursor, sortField string, sortOrder int, filtersHash string) (bson.M, error) {
	pageCursor, err := decodePageCursor(cursor)
	if err != nil {
		return nil, err
	}
	if !pageCursor.matches(sortField, sortOrder) {
		return nil, fmt.Errorf("%w: cursor does not match requested sort", ErrInvalidCursor)
	}
	if pageCursor.FiltersHash != filtersHash {
		return nil, fmt.Errorf("%w: cursor does not match current filters", ErrInvalidCursor)
	}
	return pageCursor.mongoFilter()
}

// generateCursor generates an opaque cursor from the last document in a page.
func (r *repository[T]) generateCursor(doc T, sortField string, sortOrder int, filtersHash string, pageOffset int) (string, error) {
	docBytes, err := bson.Marshal(doc)
	if err != nil {
		return "", fmt.Errorf("failed to marshal document: %w", err)
	}

	var docMap bson.M
	if err = bson.Unmarshal(docBytes, &docMap); err != nil {
		return "", fmt.Errorf("failed to unmarshal document: %w", err)
	}

	sortValue, err := documentSortValue(docMap, sortField)
	if err != nil {
		return "", err
	}

	idHex, err := documentIDHex(docMap)
	if err != nil {
		return "", err
	}

	return encodePageCursor(sortField, sortOrder, sortValue, idHex, filtersHash, pageOffset)
}

// getModelID extracts the ID from a model
func (r *repository[T]) getModelID(model *T) (string, error) {
	if idGetter, ok := any(model).(IDGetter); ok {
		return idGetter.GetID(), nil
	}

	// Fallback: try to get ID from BSON
	docBytes, err := bson.Marshal(model)
	if err != nil {
		return "", fmt.Errorf("failed to marshal model: %w", err)
	}

	var docMap bson.M
	if err = bson.Unmarshal(docBytes, &docMap); err != nil {
		return "", fmt.Errorf("failed to unmarshal model: %w", err)
	}

	if id, exists := docMap["_id"]; exists {
		if oid, ok := id.(bson.ObjectID); ok {
			return oid.Hex(), nil
		}
	}

	return "", errors.New("model does not have an ID field")
}