using Microsoft.Extensions.Logging;
using MongoDB.Bson;
using MongoDB.Driver;

namespace StellaOps.Feedser.Storage.Mongo;

public sealed class MongoSourceStateRepository : ISourceStateRepository
{
    private readonly IMongoCollection<SourceStateDocument> _collection;
    private const int MaxFailureReasonLength = 1024;

    private readonly ILogger<MongoSourceStateRepository> _logger;

    public MongoSourceStateRepository(IMongoDatabase database, ILogger<MongoSourceStateRepository> logger)
    {
        _collection = (database ?? throw new ArgumentNullException(nameof(database)))
            .GetCollection<SourceStateDocument>(MongoStorageDefaults.Collections.SourceState);
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public async Task<SourceStateRecord?> TryGetAsync(string sourceName, CancellationToken cancellationToken)
    {
        var cursor = await _collection.FindAsync(x => x.SourceName == sourceName, cancellationToken: cancellationToken).ConfigureAwait(false);
        var document = await cursor.FirstOrDefaultAsync(cancellationToken).ConfigureAwait(false);
        return document?.ToRecord();
    }

    public async Task<SourceStateRecord> UpsertAsync(SourceStateRecord record, CancellationToken cancellationToken)
    {
        var document = SourceStateDocumentExtensions.FromRecord(record with { UpdatedAt = DateTimeOffset.UtcNow });
        await _collection.ReplaceOneAsync(
            x => x.SourceName == record.SourceName,
            document,
            new ReplaceOptions { IsUpsert = true },
            cancellationToken).ConfigureAwait(false);

        _logger.LogDebug("Upserted source state for {Source}", record.SourceName);
        return document.ToRecord();
    }

    public async Task<SourceStateRecord?> UpdateCursorAsync(string sourceName, BsonDocument cursor, DateTimeOffset completedAt, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrEmpty(sourceName);
        var update = Builders<SourceStateDocument>.Update
            .Set(x => x.Cursor, cursor ?? new BsonDocument())
            .Set(x => x.LastSuccess, completedAt.UtcDateTime)
            .Set(x => x.FailCount, 0)
            .Set(x => x.BackoffUntil, (DateTime?)null)
            .Set(x => x.LastFailureReason, null)
            .Set(x => x.UpdatedAt, DateTime.UtcNow)
            .SetOnInsert(x => x.SourceName, sourceName);

        var options = new FindOneAndUpdateOptions<SourceStateDocument>
        {
            ReturnDocument = ReturnDocument.After,
            IsUpsert = true,
        };

        var document = await _collection
            .FindOneAndUpdateAsync<SourceStateDocument, SourceStateDocument>(
                x => x.SourceName == sourceName,
                update,
                options,
                cancellationToken)
            .ConfigureAwait(false);
        return document?.ToRecord();
    }

    public async Task<SourceStateRecord?> MarkFailureAsync(string sourceName, DateTimeOffset failedAt, TimeSpan? backoff, string? failureReason, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrEmpty(sourceName);
        var reasonValue = NormalizeFailureReason(failureReason);
        var update = Builders<SourceStateDocument>.Update
            .Inc(x => x.FailCount, 1)
            .Set(x => x.LastFailure, failedAt.UtcDateTime)
            .Set(x => x.BackoffUntil, backoff.HasValue ? failedAt.UtcDateTime.Add(backoff.Value) : null)
            .Set(x => x.LastFailureReason, reasonValue)
            .Set(x => x.UpdatedAt, DateTime.UtcNow)
            .SetOnInsert(x => x.SourceName, sourceName);

        var options = new FindOneAndUpdateOptions<SourceStateDocument>
        {
            ReturnDocument = ReturnDocument.After,
            IsUpsert = true,
        };

        var document = await _collection
            .FindOneAndUpdateAsync<SourceStateDocument, SourceStateDocument>(
                x => x.SourceName == sourceName,
                update,
                options,
                cancellationToken)
            .ConfigureAwait(false);
        return document?.ToRecord();
    }

    private static string? NormalizeFailureReason(string? reason)
    {
        if (string.IsNullOrWhiteSpace(reason))
        {
            return null;
        }

        var trimmed = reason.Trim();
        if (trimmed.Length <= MaxFailureReasonLength)
        {
            return trimmed;
        }

        return trimmed[..MaxFailureReasonLength];
    }
}
