// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;
using MongoDB.Bson;

namespace Microsoft.SemanticKernel.Connectors.MongoDB;

/// <summary>
/// Contains mapping helpers to use when searching for documents using MongoDB.
/// </summary>
internal static class MongoCollectionSearchMapping
{
    /// <summary>Returns search part of the search query.</summary>
    public static BsonDocument GetSearchQuery<TVector>(
        TVector vector,
        string indexName,
        string vectorPropertyName,
        int limit,
        int numCandidates,
        BsonDocument? filter)
    {
        // Docs: https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage
        var searchQuery = new BsonDocument
        {
            { "index", indexName },
            { "queryVector", BsonArray.Create(vector) },
            { "path", vectorPropertyName },
            { "limit", limit },
            { "numCandidates", numCandidates },
        };

        if (filter is not null)
        {
            searchQuery["filter"] = filter;
        }

        return new BsonDocument
        {
            { "$vectorSearch", searchQuery }
        };
    }

    /// <summary>Returns projection part of the search query to return similarity score together with document.</summary>
    public static BsonDocument GetProjectionQuery(string scorePropertyName, string documentPropertyName)
    {
        return new BsonDocument
        {
            {
                "$project", new BsonDocument
                {
                    { scorePropertyName, new BsonDocument { { "$meta", "vectorSearchScore" } } },
                    { documentPropertyName, "$$ROOT" }
                }
            }
        };
    }

    /// <summary>Returns a $match stage to filter results by score threshold.</summary>
    /// <remarks>
    /// MongoDB Atlas Vector Search returns a similarity score where higher values mean more similar,
    /// so we filter with $gte to keep results at or above the threshold.
    /// </remarks>
    public static BsonDocument GetScoreThresholdMatchQuery(string scorePropertyName, double scoreThreshold)
        => new()
        {
            {
                "$match", new BsonDocument
                {
                    { scorePropertyName, new BsonDocument { { "$gte", scoreThreshold } } }
                }
            }
        };

    /// <summary>Returns a pipeline for hybrid search using vector search and full text search.</summary>
    public static BsonDocument[] GetHybridSearchPipeline<TVector>(
        TVector vector,
        ICollection<string> keywords,
        string collectionName,
        string vectorIndexName,
        string fullTextSearchIndexName,
        string vectorPropertyName,
        string textPropertyName,
        string scorePropertyName,
        string documentPropertyName,
        int limit,
        int numCandidates,
        BsonDocument? filter)
    {
        // Create the FullTextSearch pipeline first.
        var ftsPipeline = new List<BsonDocument>
        {
            // The full text search stage.
            GetFullTextSearchQuery(keywords, fullTextSearchIndexName, textPropertyName, filter),
            // Limit the results to the maximum that we may require.
            new()
            {
                {
                    "$limit", limit
                }
            },
            // Converts the list of documents to a single document with an array property containing all the source documents.
            GroupDocsSection(),
            // Creates separate documents again where each has a new rank property based on the index of the document.
            UnwindDocsArraySection(),
            // Add a weighted score based on the rank of the document.
            AddScore("fts_score", 0.9),
            // Project the score, the id and the original document as properties, so that we can join with the vector search results on id.
            ProjectWithScore("fts_score"),
        };

        // Add filtering to the FullTextSearch pipeline if filter is provided.
        if (filter is not null)
        {
            // Insert filter at the second position, since
            // MongoDB requires search to be the first stage.
            ftsPipeline.Insert(1, new BsonDocument
            {
                {
                    "$match", filter
                }
            });
        }

        // Create combined pipeline with the vector search part first.
        var pipeline = new BsonDocument[]
        {
            // The vector search stage.
            GetSearchQuery(vector, vectorIndexName, vectorPropertyName, limit, numCandidates, filter),
            // Converts the list of documents to a single document with an array property containing all the source documents.
            GroupDocsSection(),
            // Creates separate documents again where each has a new rank property based on the index of the document.
            UnwindDocsArraySection(),
            // Add a weighted score based on the rank of the document.
            AddScore("vs_score", 0.1),
            // Project the score, the id and the original document as properties, so that we can join with the vector search results on id.
            ProjectWithScore("vs_score"),
            // Union the vector search results with the results from the full text search pipeilne.
            new()
            {
                {
                    "$unionWith", new BsonDocument
                    {
                        { "coll", collectionName },
                        { "pipeline", new BsonArray(ftsPipeline) }
                    }
                }
            },
            // Group by id and store scores from each pipeline, so that we don't have duplicate documents.
            new()
            {
                {
                    "$group", new BsonDocument
                    {
                        { "_id", "$_id" },
                        { "docs", new BsonDocument { { "$first", "$docs" } } },
                        { "vs_score", new BsonDocument { { "$max", "$vs_score" } } },
                        { "fts_score", new BsonDocument { { "$max", "$fts_score" } } }
                    }
                }
            },
            // If a score is missing (i.e. the document was only found in the other pipeline), default the missing score to 0.
            new()
            {
                {
                    "$project", new BsonDocument
                    {
                        { "_id", 1 },
                        { "docs", 1 },
                        { "vs_score", new BsonDocument { { "$ifNull", new BsonArray { "$vs_score", 0 } } } },
                        { "fts_score", new BsonDocument { { "$ifNull", new BsonArray { "$fts_score", 0 } } } }
                    }
                }
            },
            // Calculate a combined score based on the vector search and full text search scores.
            new()
            {
                {
                    "$project", new BsonDocument
                    {
                        { scorePropertyName, new BsonDocument { { "$add", new BsonArray { "$fts_score", "$vs_score" } } } },
                        { "vs_score", 1 },
                        { "fts_score", 1 },
                        { documentPropertyName, "$docs" }
                    }
                }
            },
            // Sort by score desc.
            new()
            {
                {
                    "$sort", new BsonDocument
                    {
                        { scorePropertyName, -1 }
                    }
                }
            },
            // Take the required N results.
            new()
            {
                {
                    "$limit", limit
                }
            },
        };

        return pipeline;
    }

    /// <summary>Builds the full text search query stage.</summary>
    private static BsonDocument GetFullTextSearchQuery(
        ICollection<string> keywords,
        string fullTextSearchIndexName,
        string textPropertyName,
        BsonDocument? filter)
    {
        var fullTextSearchQuery = new BsonDocument
        {
            {
                "$search", new BsonDocument
                {
                    { "index", fullTextSearchIndexName },
                    { "text",
                        new BsonDocument
                        {
                            { "query", new BsonArray(keywords) },
                            { "path", textPropertyName },
                            { "matchCriteria", "any" }
                        }
                    }
                }
            }
        };

        return fullTextSearchQuery;
    }

    /// <summary>Create a stage that groups all documents into a single document with an array property containing all the source documents.</summary>
    private static BsonDocument GroupDocsSection()
    {
        return new BsonDocument
        {
            {
                "$group", new BsonDocument
                {
                    { "_id", BsonNull.Value },
                    { "docs", new BsonDocument { { "$push", "$$ROOT" } } }
                }
            }
        };
    }

    /// <summary>Creates a stage that splits an array of documents from a single document into separate documents and adds an index property for each document based on its index in the array.</summary>
    private static BsonDocument UnwindDocsArraySection()
    {
        return new BsonDocument
        {
            {
                "$unwind", new BsonDocument
                {
                    { "path", "$docs" },
                    { "includeArrayIndex", "rank" }
                }
            }
        };
    }

    /// <summary>Adds a weighted score to each document based on the rank property on the document.</summary>
    private static BsonDocument AddScore(string scoreName, double weight)
    {
        return new()
        {
            {
                "$addFields", new BsonDocument
                {
                    {
                        scoreName, new BsonDocument
                        {
                            {
                                "$multiply", new BsonArray()
                                {
                                    weight,
                                    new BsonDocument
                                    {
                                        { "$divide", new BsonArray()
                                            {
                                                1.0,
                                                new BsonDocument
                                                {
                                                    { "$add", new BsonArray()
                                                        {
                                                            "$rank",
                                                            60
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    },
                }
            }
        };
    }

    /// <summary>Projects the score, the id and the original document as properties.</summary>
    private static BsonDocument ProjectWithScore(string scoreName)
    {
        return new()
        {
            {
                "$project", new BsonDocument
                {
                    { scoreName, 1 },
                    { "_id", "$docs._id" },
                    { "docs", "$docs" }
                }
            }
        };
    }
}
