namespace Extensions
{
    using System;

    using GreenPipes;

    using MassTransit;
    using MassTransit.QuartzIntegration;
    using MassTransit.QuartzIntegration.Configuration;
    using MassTransit.Scheduling;

    using Quartz;

    public static class InMemoryBusFactoryConfiguratorExtensions
    {
        public static void UseTestInMemoryScheduler(this IBusFactoryConfigurator configurator, IScheduler scheduler, string queueName = "quartz")
        {
            configurator.ReceiveEndpoint(queueName, e =>
            {
                var partitioner = configurator.CreatePartitioner(16);

                e.Consumer(() => new ScheduleMessageConsumer(scheduler), x =>
                    x.Message<ScheduleMessage>(m => m.UsePartitioner(partitioner, p => p.Message.CorrelationId)));

                e.Consumer(() => new CancelScheduledMessageConsumer(scheduler), x =>
                    x.Message<CancelScheduledMessage>(m => m.UsePartitioner(partitioner, p => p.Message.TokenId)));

                configurator.UseMessageScheduler(e.InputAddress);

                var observer = new SchedulerBusObserver(scheduler, e.InputAddress);
                configurator.ConnectBusObserver(observer);
            });
        }
    }
}