using MongoDB.Bson;
using MongoDB.Driver;
using StellaOps.Scanner.Storage.Catalog;
using StellaOps.Scanner.Storage.Mongo;

namespace StellaOps.Scanner.Storage.Repositories;

public sealed class JobRepository
{
    private readonly MongoCollectionProvider _collections;
    private readonly TimeProvider _timeProvider;

    public JobRepository(MongoCollectionProvider collections, TimeProvider? timeProvider = null)
    {
        _collections = collections ?? throw new ArgumentNullException(nameof(collections));
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    public async Task<JobDocument> InsertAsync(JobDocument document, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(document);
        document.CreatedAtUtc = _timeProvider.GetUtcNow().UtcDateTime;
        await _collections.Jobs.InsertOneAsync(document, cancellationToken: cancellationToken).ConfigureAwait(false);
        return document;
    }

    public async Task<bool> TryTransitionAsync(string jobId, JobState expected, JobState next, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);
        var now = _timeProvider.GetUtcNow().UtcDateTime;
        var update = Builders<JobDocument>.Update
            .Set(x => x.State, next)
            .Set(x => x.HeartbeatAtUtc, now);

        if (next == JobState.Running)
        {
            update = update.Set(x => x.StartedAtUtc, now);
        }

        if (next is JobState.Succeeded or JobState.Failed or JobState.Cancelled)
        {
            update = update.Set(x => x.CompletedAtUtc, now);
        }

        var result = await _collections.Jobs.UpdateOneAsync(
            Builders<JobDocument>.Filter.And(
                Builders<JobDocument>.Filter.Eq(x => x.Id, jobId),
                Builders<JobDocument>.Filter.Eq(x => x.State, expected)),
            update,
            cancellationToken: cancellationToken).ConfigureAwait(false);

        return result.ModifiedCount == 1;
    }

    public async Task<JobDocument?> GetAsync(string jobId, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);
        return await _collections.Jobs
            .Find(x => x.Id == jobId)
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    public Task<List<JobDocument>> ListStaleAsync(TimeSpan heartbeatThreshold, CancellationToken cancellationToken)
    {
        if (heartbeatThreshold <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(nameof(heartbeatThreshold));
        }

        var cutoff = _timeProvider.GetUtcNow().UtcDateTime - heartbeatThreshold;
        var filter = Builders<JobDocument>.Filter.And(
            Builders<JobDocument>.Filter.Eq(x => x.State, JobState.Running),
            Builders<JobDocument>.Filter.Lt(x => x.HeartbeatAtUtc, cutoff));

        return _collections.Jobs.Find(filter).ToListAsync(cancellationToken);
    }
}
