using Microsoft.Extensions.Logging;
using MongoDB.Driver;
using StellaOps.Feedser.Core.Jobs;

namespace StellaOps.Feedser.Storage.Mongo;

public sealed class MongoLeaseStore : ILeaseStore
{
    private readonly IMongoCollection<JobLeaseDocument> _collection;
    private readonly ILogger<MongoLeaseStore> _logger;

    public MongoLeaseStore(IMongoCollection<JobLeaseDocument> collection, ILogger<MongoLeaseStore> logger)
    {
        _collection = collection ?? throw new ArgumentNullException(nameof(collection));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public async Task<JobLease?> TryAcquireAsync(string key, string holder, TimeSpan leaseDuration, DateTimeOffset now, CancellationToken cancellationToken)
    {
        var nowUtc = now.UtcDateTime;
        var ttlUtc = nowUtc.Add(leaseDuration);

        var filter = Builders<JobLeaseDocument>.Filter.Eq(x => x.Key, key)
            & Builders<JobLeaseDocument>.Filter.Or(
                Builders<JobLeaseDocument>.Filter.Lte(x => x.TtlAt, nowUtc),
                Builders<JobLeaseDocument>.Filter.Eq(x => x.Holder, holder));

        var update = Builders<JobLeaseDocument>.Update
            .Set(x => x.Holder, holder)
            .Set(x => x.AcquiredAt, nowUtc)
            .Set(x => x.HeartbeatAt, nowUtc)
            .Set(x => x.LeaseMs, (long)leaseDuration.TotalMilliseconds)
            .Set(x => x.TtlAt, ttlUtc);

        var options = new FindOneAndUpdateOptions<JobLeaseDocument>
        {
            ReturnDocument = ReturnDocument.After,
        };

        var updated = await _collection.FindOneAndUpdateAsync(filter, update, options, cancellationToken).ConfigureAwait(false);
        if (updated is not null)
        {
            _logger.LogDebug("Lease {Key} acquired by {Holder}", key, holder);
            return updated.ToLease();
        }

        try
        {
            var document = new JobLeaseDocument
            {
                Key = key,
                Holder = holder,
                AcquiredAt = nowUtc,
                HeartbeatAt = nowUtc,
                LeaseMs = (long)leaseDuration.TotalMilliseconds,
                TtlAt = ttlUtc,
            };

            await _collection.InsertOneAsync(document, cancellationToken: cancellationToken).ConfigureAwait(false);
            _logger.LogDebug("Lease {Key} inserted for {Holder}", key, holder);
            return document.ToLease();
        }
        catch (MongoWriteException ex) when (ex.WriteError.Category == ServerErrorCategory.DuplicateKey)
        {
            _logger.LogDebug(ex, "Lease {Key} already held by another process", key);
            return null;
        }
    }

    public async Task<JobLease?> HeartbeatAsync(string key, string holder, TimeSpan leaseDuration, DateTimeOffset now, CancellationToken cancellationToken)
    {
        var nowUtc = now.UtcDateTime;
        var ttlUtc = nowUtc.Add(leaseDuration);

        var filter = Builders<JobLeaseDocument>.Filter.Eq(x => x.Key, key)
            & Builders<JobLeaseDocument>.Filter.Eq(x => x.Holder, holder);

        var update = Builders<JobLeaseDocument>.Update
            .Set(x => x.HeartbeatAt, nowUtc)
            .Set(x => x.LeaseMs, (long)leaseDuration.TotalMilliseconds)
            .Set(x => x.TtlAt, ttlUtc);

        var updated = await _collection.FindOneAndUpdateAsync(
            filter,
            update,
            new FindOneAndUpdateOptions<JobLeaseDocument>
            {
                ReturnDocument = ReturnDocument.After,
            },
            cancellationToken).ConfigureAwait(false);

        if (updated is null)
        {
            _logger.LogDebug("Heartbeat rejected for lease {Key} held by {Holder}", key, holder);
        }

        return updated?.ToLease();
    }

    public async Task<bool> ReleaseAsync(string key, string holder, CancellationToken cancellationToken)
    {
        var result = await _collection.DeleteOneAsync(
            Builders<JobLeaseDocument>.Filter.Eq(x => x.Key, key)
            & Builders<JobLeaseDocument>.Filter.Eq(x => x.Holder, holder),
            cancellationToken).ConfigureAwait(false);

        if (result.DeletedCount == 0)
        {
            _logger.LogDebug("Lease {Key} not released by {Holder}; no matching document", key, holder);
            return false;
        }

        _logger.LogDebug("Lease {Key} released by {Holder}", key, holder);
        return true;
    }
}
