step3/api/QuestionsApi/Services/AnswerStatisticsStore.cs

using MongoDB.Bson;
using MongoDB.Driver;
using QuestionsApi.Models;
using System.Collections.Concurrent;

namespace QuestionsApi.Services;

public sealed class AnswerStatisticsStore
{
    private const string StateId = "answer-statistics-v2";
    private static readonly TimeSpan CacheDuration = TimeSpan.FromMinutes(5);

    private readonly IMongoCollection<QuestionAttemptCounter> _counters;
    private readonly IMongoCollection<AnswerStatisticsState> _state;
    private readonly IMongoCollection<QuestionPresentationBaseline> _baselines;
    private readonly QuestionAttemptStore _attempts;
    private readonly QuestionStore _questions;
    private readonly ConcurrentDictionary<string, CachedPercentages> _cache = new();
    private volatile bool _backfillComplete;

    public AnswerStatisticsStore(
        MongoDatabaseProvider mongo,
        QuestionAttemptStore attempts,
        QuestionStore questions)
    {
        _attempts = attempts;
        _questions = questions;
        _counters = mongo.Database.GetCollection<QuestionAttemptCounter>("QuestionChoiceExposureCounters");
        _state = mongo.Database.GetCollection<AnswerStatisticsState>("SystemState");
        _baselines = mongo.Database.GetCollection<QuestionPresentationBaseline>("QuestionPresentationBaselines");

        var counterIndex = Builders<QuestionAttemptCounter>.IndexKeys
            .Ascending(counter => counter.QuestionExternalId)
            .Ascending(counter => counter.Bucket)
            .Ascending(counter => counter.ChoiceId);
        _counters.Indexes.CreateOne(new CreateIndexModel<QuestionAttemptCounter>(
            counterIndex,
            new CreateIndexOptions { Unique = true }));
    }

    public async Task IncrementAsync(QuestionAttempt attempt)
    {
        if (attempt.IsSkip || !attempt.TracksPresentation || attempt.PresentedChoiceIds.Count == 0)
            return;

        var now = DateTime.UtcNow;
        var writes = attempt.PresentedChoiceIds.Distinct(StringComparer.Ordinal).Select(choiceId =>
        {
            var filter = CounterFilter(attempt.QuestionExternalId, attempt.CounterBucket, choiceId);
            var update = Builders<QuestionAttemptCounter>.Update
                .Inc(counter => counter.PresentationCount, 1)
                .Inc(counter => counter.SelectionCount, choiceId == attempt.ChoiceId ? 1 : 0)
                .Set(counter => counter.UpdatedUtc, now)
                .SetOnInsert(counter => counter.QuestionExternalId, attempt.QuestionExternalId)
                .SetOnInsert(counter => counter.Bucket, attempt.CounterBucket)
                .SetOnInsert(counter => counter.ChoiceId, choiceId);
            return new UpdateOneModel<QuestionAttemptCounter>(filter, update) { IsUpsert = true };
        }).ToList();

        try
        {
            await _counters.BulkWriteAsync(writes, new BulkWriteOptions { IsOrdered = false });
        }
        catch (MongoBulkWriteException<QuestionAttemptCounter> exception)
            when (exception.WriteErrors.All(error => error.Category == ServerErrorCategory.DuplicateKey))
        {
            await _counters.BulkWriteAsync(writes, new BulkWriteOptions { IsOrdered = false });
        }
        _cache.TryRemove(attempt.QuestionExternalId, out _);
    }

    public async Task<Dictionary<string, int>> GetPercentagesForQuestionAsync(string questionExternalId)
    {
        var result = await GetPercentagesByQuestionAsync([questionExternalId]);
        return result.GetValueOrDefault(questionExternalId, []);
    }

    public async Task<Dictionary<string, Dictionary<string, int>>> GetPercentagesByQuestionAsync(
        IEnumerable<string> questionExternalIds)
    {
        var ids = questionExternalIds.Where(id => !string.IsNullOrWhiteSpace(id)).Distinct().ToList();
        var output = new Dictionary<string, Dictionary<string, int>>();
        var missing = new List<string>();
        var now = DateTime.UtcNow;

        foreach (var id in ids)
        {
            if (_cache.TryGetValue(id, out var cached) && cached.ExpiresUtc > now)
                output[id] = new Dictionary<string, int>(cached.Percentages);
            else
                missing.Add(id);
        }
        if (missing.Count == 0 || !await IsBackfillCompleteAsync())
            return output;

        var totals = await _counters.Aggregate()
            .Match(counter => missing.Contains(counter.QuestionExternalId))
            .Group(
                counter => new { counter.QuestionExternalId, counter.ChoiceId },
                group => new QuestionAttemptCounterTotal
                {
                    Key = new QuestionChoiceCounterKey
                    {
                        QuestionExternalId = group.Key.QuestionExternalId,
                        ChoiceId = group.Key.ChoiceId
                    },
                    SelectionCount = group.Sum(counter => counter.SelectionCount),
                    PresentationCount = group.Sum(counter => counter.PresentationCount)
                })
            .ToListAsync();

        var grouped = totals.GroupBy(total => total.Key.QuestionExternalId).ToDictionary(group => group.Key);
        foreach (var id in missing)
        {
            var percentages = grouped.TryGetValue(id, out var choices)
                ? choices.Where(item => item.PresentationCount > 0).ToDictionary(
                    item => item.Key.ChoiceId,
                    item => (int)Math.Round(item.SelectionCount * 100.0 / item.PresentationCount))
                : [];
            output[id] = percentages;
            _cache[id] = new CachedPercentages(percentages, now.Add(CacheDuration));
        }
        return output;
    }

    public async Task<bool> TryAcquireBackfillLeaseAsync(string owner, TimeSpan duration)
    {
        var now = DateTime.UtcNow;
        var filter = Builders<AnswerStatisticsState>.Filter.And(
            Builders<AnswerStatisticsState>.Filter.Eq(state => state.Id, StateId),
            Builders<AnswerStatisticsState>.Filter.Eq(state => state.Completed, false),
            Builders<AnswerStatisticsState>.Filter.Or(
                Builders<AnswerStatisticsState>.Filter.Lt(state => state.LeaseUntilUtc, now),
                Builders<AnswerStatisticsState>.Filter.Eq(state => state.LeaseOwner, owner)));
        var update = Builders<AnswerStatisticsState>.Update
            .SetOnInsert(state => state.Id, StateId)
            .SetOnInsert(state => state.Completed, false)
            .SetOnInsert(state => state.CutoverUtc, now)
            .Set(state => state.LeaseOwner, owner)
            .Set(state => state.LeaseUntilUtc, now.Add(duration))
            .Set(state => state.UpdatedUtc, now);

        try
        {
            var leased = await _state.FindOneAndUpdateAsync(filter, update,
                new FindOneAndUpdateOptions<AnswerStatisticsState>
                {
                    IsUpsert = true,
                    ReturnDocument = ReturnDocument.After
                });
            return leased?.LeaseOwner == owner;
        }
        catch (MongoWriteException exception) when (exception.WriteError.Category == ServerErrorCategory.DuplicateKey)
        {
            return false;
        }
        catch (MongoCommandException exception) when (exception.Code == 11000)
        {
            return false;
        }
    }

    public async Task MarkBackfillCompleteAsync(string owner)
    {
        await _state.UpdateOneAsync(
            state => state.Id == StateId && state.LeaseOwner == owner,
            Builders<AnswerStatisticsState>.Update
                .Set(state => state.Completed, true)
                .Set(state => state.UpdatedUtc, DateTime.UtcNow));
        _backfillComplete = true;
    }

    public Task RebuildAllAsync() => RebuildAsync(null);

    public Task RebuildQuestionsAsync(IReadOnlyCollection<string> questionExternalIds) =>
        questionExternalIds.Count == 0 ? Task.CompletedTask : RebuildAsync(questionExternalIds);

    public async Task<bool> IsBackfillCompleteAsync()
    {
        if (_backfillComplete)
            return true;
        _backfillComplete = (await _state.Find(state => state.Id == StateId).FirstOrDefaultAsync())?.Completed == true;
        return _backfillComplete;
    }

    private async Task RebuildAsync(IReadOnlyCollection<string>? questionExternalIds)
    {
        await EnsureBaselinesAsync();
        var state = await _state.Find(item => item.Id == StateId).FirstAsync();
        var baselines = await _baselines.Find(Builders<QuestionPresentationBaseline>.Filter.Empty).ToListAsync();
        var baselineByQuestion = baselines.ToDictionary(item => item.QuestionExternalId, item => item.ChoiceIds);
        var match = new BsonDocument
        {
            { "IsSkip", false },
            { "$or", new BsonArray
                {
                    new BsonDocument("TracksPresentation", true),
                    new BsonDocument("CreatedUtc", new BsonDocument("$lt", state.CutoverUtc))
                }
            }
        };
        if (questionExternalIds is not null)
            match["QuestionExternalId"] = new BsonDocument("$in", new BsonArray(questionExternalIds));

        var selectionPipeline = new[]
        {
            new BsonDocument("$match", match),
            new BsonDocument("$group", new BsonDocument
            {
                { "_id", new BsonDocument
                    {
                        { "q", "$QuestionExternalId" },
                        { "b", new BsonDocument("$ifNull", new BsonArray { "$CounterBucket", 0 }) },
                        { "c", "$ChoiceId" }
                    }
                },
                { "count", new BsonDocument("$sum", 1) }
            })
        };
        var presentationMatch = new BsonDocument(match) { ["TracksPresentation"] = true };
        var presentationPipeline = new[]
        {
            new BsonDocument("$match", presentationMatch),
            new BsonDocument("$unwind", "$PresentedChoiceIds"),
            new BsonDocument("$group", new BsonDocument
            {
                { "_id", new BsonDocument
                    {
                        { "q", "$QuestionExternalId" },
                        { "b", new BsonDocument("$ifNull", new BsonArray { "$CounterBucket", 0 }) },
                        { "c", "$PresentedChoiceIds" }
                    }
                },
                { "count", new BsonDocument("$sum", 1) }
            })
        };
        var legacyMatch = new BsonDocument
        {
            { "IsSkip", false },
            { "TracksPresentation", new BsonDocument("$ne", true) },
            { "CreatedUtc", new BsonDocument("$lt", state.CutoverUtc) }
        };
        if (questionExternalIds is not null)
            legacyMatch["QuestionExternalId"] = new BsonDocument("$in", new BsonArray(questionExternalIds));
        var legacyPipeline = new[]
        {
            new BsonDocument("$match", legacyMatch),
            new BsonDocument("$group", new BsonDocument
            {
                { "_id", new BsonDocument
                    {
                        { "q", "$QuestionExternalId" },
                        { "b", new BsonDocument("$ifNull", new BsonArray { "$CounterBucket", 0 }) }
                    }
                },
                { "count", new BsonDocument("$sum", 1) }
            })
        };

        var selections = await _attempts.Collection.Aggregate<BsonDocument>(selectionPipeline).ToListAsync();
        var presentations = await _attempts.Collection.Aggregate<BsonDocument>(presentationPipeline).ToListAsync();
        var legacy = await _attempts.Collection.Aggregate<BsonDocument>(legacyPipeline).ToListAsync();
        var totals = new Dictionary<(string Question, int Bucket, string Choice), (long Selected, long Presented)>();

        foreach (var row in selections)
        {
            var key = CounterKey(row["_id"].AsBsonDocument);
            totals[key] = (row["count"].ToInt64(), totals.GetValueOrDefault(key).Presented);
        }
        foreach (var row in presentations)
        {
            var key = CounterKey(row["_id"].AsBsonDocument);
            totals[key] = (totals.GetValueOrDefault(key).Selected, row["count"].ToInt64());
        }
        foreach (var row in legacy)
        {
            var id = row["_id"].AsBsonDocument;
            var questionId = id["q"].AsString;
            if (!baselineByQuestion.TryGetValue(questionId, out var choiceIds))
                continue;
            foreach (var choiceId in choiceIds)
            {
                var key = (questionId, id["b"].ToInt32(), choiceId);
                totals[key] = (totals.GetValueOrDefault(key).Selected,
                    totals.GetValueOrDefault(key).Presented + row["count"].ToInt64());
            }
        }

        var now = DateTime.UtcNow;
        var writes = totals.Select(item =>
            new ReplaceOneModel<QuestionAttemptCounter>(
                CounterFilter(item.Key.Question, item.Key.Bucket, item.Key.Choice),
                new QuestionAttemptCounter
                {
                    QuestionExternalId = item.Key.Question,
                    Bucket = item.Key.Bucket,
                    ChoiceId = item.Key.Choice,
                    SelectionCount = item.Value.Selected,
                    PresentationCount = item.Value.Presented,
                    UpdatedUtc = now
                })
            {
                IsUpsert = true
            }).ToList();
        if (writes.Count > 0)
            await _counters.BulkWriteAsync(writes, new BulkWriteOptions { IsOrdered = false });

        if (questionExternalIds is null)
            _cache.Clear();
        else
            foreach (var id in questionExternalIds)
                _cache.TryRemove(id, out _);
    }

    private async Task EnsureBaselinesAsync()
    {
        var questions = await _questions.GetAllAsync();
        var writes = questions.Select(question =>
            new UpdateOneModel<QuestionPresentationBaseline>(
                Builders<QuestionPresentationBaseline>.Filter.Eq(item => item.QuestionExternalId, question.ExternalId),
                Builders<QuestionPresentationBaseline>.Update
                    .SetOnInsert(item => item.QuestionExternalId, question.ExternalId)
                    .SetOnInsert(item => item.ChoiceIds, question.Choices.Select(choice => choice.Id).ToList())
                    .SetOnInsert(item => item.CreatedUtc, DateTime.UtcNow))
            {
                IsUpsert = true
            }).ToList();
        if (writes.Count > 0)
            await _baselines.BulkWriteAsync(writes, new BulkWriteOptions { IsOrdered = false });
    }

    private static FilterDefinition<QuestionAttemptCounter> CounterFilter(string questionId, int bucket, string choiceId) =>
        Builders<QuestionAttemptCounter>.Filter.And(
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.QuestionExternalId, questionId),
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.Bucket, bucket),
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.ChoiceId, choiceId));

    private static (string Question, int Bucket, string Choice) CounterKey(BsonDocument id) =>
        (id["q"].AsString, id["b"].ToInt32(), id["c"].AsString);

    private sealed record CachedPercentages(Dictionary<string, int> Percentages, DateTime ExpiresUtc);
}