package com.atlassian.scheduler.caesium.impl;
   
   import com.atlassian.annotations.Internal;
   import com.atlassian.scheduler.JobRunner;
   import com.atlassian.scheduler.JobRunnerRequest;
   import com.atlassian.scheduler.JobRunnerResponse;
   import com.atlassian.scheduler.SchedulerRuntimeException;
   import com.atlassian.scheduler.SchedulerServiceException;
   import com.atlassian.scheduler.caesium.migration.LazyMigratingParameterMapSerializer;
  import com.atlassian.scheduler.caesium.spi.CaesiumSchedulerConfiguration;
  import com.atlassian.scheduler.caesium.spi.ClusteredJob;
  import com.atlassian.scheduler.caesium.spi.ClusteredJobDao;
  import com.atlassian.scheduler.config.CronScheduleInfo;
  import com.atlassian.scheduler.config.IntervalScheduleInfo;
  import com.atlassian.scheduler.config.JobConfig;
  import com.atlassian.scheduler.config.JobId;
  import com.atlassian.scheduler.config.JobRunnerKey;
  import com.atlassian.scheduler.config.RunMode;
  import com.atlassian.scheduler.config.Schedule;
  import com.atlassian.scheduler.core.AbstractSchedulerService;
  import com.atlassian.scheduler.core.JobLauncher;
  import com.atlassian.scheduler.core.spi.RunDetailsDao;
  import com.atlassian.scheduler.core.status.LazyJobDetails;
  import com.atlassian.scheduler.core.status.SimpleJobDetails;
  import com.atlassian.scheduler.core.util.ParameterMapSerializer;
  import com.atlassian.scheduler.cron.CronSyntaxException;
  import com.atlassian.scheduler.status.JobDetails;
  import com.atlassian.util.concurrent.Sink;
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import javax.annotation.Nonnull;
  import javax.annotation.Nullable;
  import java.io.Serializable;
  import java.util.Collection;
  import java.util.Date;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.TreeMap;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.ThreadFactory;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import static com.atlassian.scheduler.config.RunMode.RUN_LOCALLY;
  import static com.atlassian.scheduler.config.RunMode.RUN_ONCE_PER_CLUSTER;
  import static com.atlassian.scheduler.core.util.CronExpressionQuantizer.quantizeSecondsField;
  import static com.atlassian.scheduler.core.util.TimeIntervalQuantizer.quantizeToMinutes;
  import static com.atlassian.util.concurrent.Assertions.notNull;
  
  /**
   * Simple direct implementation of a scheduler service
   */
  public class CaesiumSchedulerService extends AbstractSchedulerService {
      private static final Logger LOG = LoggerFactory.getLogger(CaesiumSchedulerService.class);
      private static final JobId REFRESH_JOB_ID = JobId.of("CaesiumSchedulerService.RefreshJob");
      private static final JobRunnerKey REFRESH_JOB_RUNNER_KEY = JobRunnerKey.of("CaesiumSchedulerService.RefreshJob");
      private static final int DEFAULT_JOB_MAP_SIZE = 256;
  
      /**
       * Used to cap the number of attempts to alter the state of a clustered job.
       */
      private static final int MAX_TRIES = 50;
  
      /**
       * The default number of workers to use when the configuration gives a non-positive value
       */
      private static final int DEFAULT_WORKER_COUNT = 4;
  
      private final ConcurrentMap<JobId, JobDetails> localJobs = new ConcurrentHashMap<JobId, JobDetails>(DEFAULT_JOB_MAP_SIZE);
      private final RefreshJob refreshJob = new RefreshJob();
      private final AtomicBoolean started = new AtomicBoolean();
  
      private final ClusteredJobDao clusteredJobDao;
      private final CaesiumSchedulerConfiguration config;
      private final SchedulerQueue queue;
      private final RunTimeCalculator runTimeCalculator;
  
      @SuppressWarnings("unused")
      public CaesiumSchedulerService(final CaesiumSchedulerConfiguration config, final RunDetailsDao runDetailsDao,
                                     final ClusteredJobDao clusteredJobDao) {
          this(config, runDetailsDao, clusteredJobDao, createParameterMapSerializer(config));
      }
  
      @SuppressWarnings("unused")
      public CaesiumSchedulerService(final CaesiumSchedulerConfiguration config, final RunDetailsDao runDetailsDao,
                                     final ClusteredJobDao clusteredJobDao, final ParameterMapSerializer serializer) {
          super(runDetailsDao, notNull("serializer", serializer));
          this.clusteredJobDao = clusteredJobDao;
          this.config = config;
          this.queue = new SchedulerQueueImpl(clusteredJobDao);
          this.runTimeCalculator = new RunTimeCalculator(config);
      }
 
     @VisibleForTesting
     CaesiumSchedulerService(final CaesiumSchedulerConfiguration config,
                             final RunDetailsDao runDetailsDao, final ClusteredJobDao clusteredJobDao, final SchedulerQueue queue,
                             final RunTimeCalculator runTimeCalculator) {
         super(runDetailsDao, createParameterMapSerializer(config));
         this.clusteredJobDao = clusteredJobDao;
         this.config = config;
         this.queue = queue;
         this.runTimeCalculator = runTimeCalculator;
     }
 
     @Override
     public void scheduleJob(final JobId jobId, final JobConfig jobConfig)
             throws SchedulerServiceException {
         notNull("jobId", jobId);
         notNull("jobConfig", jobConfig);
 
         try {
             LOG.debug("scheduleJob: {}: {}", jobId, jobConfig);
 
             switch (jobConfig.getRunMode()) {
                 case RUN_LOCALLY:
                     scheduleLocalJob(jobId, jobConfig);
                     break;
                 case RUN_ONCE_PER_CLUSTER:
                     scheduleClusteredJob(jobId, jobConfig);
                     break;
                 default:
                     throw new IllegalArgumentException("Unsupported run mode: " + jobConfig.getRunMode());
             }
         } catch (SchedulerRuntimeException sre) {
             throw checked(sre);
         }
     }
 
     private void scheduleLocalJob(final JobId jobId, final JobConfig jobConfig) throws SchedulerServiceException {
         final Date nextRunTime = runTimeCalculator.firstRunTime(jobId, jobConfig);
         final Map<String, Serializable> parameters = jobConfig.getParameters();
         final SimpleJobDetails jobDetails = new SimpleJobDetails(jobId, jobConfig.getJobRunnerKey(),
                 RUN_LOCALLY, quantize(jobConfig.getSchedule()), nextRunTime,
                 getParameterMapSerializer().serializeParameters(parameters), parameters);
 
         clusteredJobDao.delete(jobId);
         localJobs.put(jobId, jobDetails);
         enqueueJob(jobId, nextRunTime);
     }
 
     private void scheduleClusteredJob(final JobId jobId, final JobConfig jobConfig) throws SchedulerServiceException {
         final Date nextRunTime = runTimeCalculator.firstRunTime(jobId, jobConfig);
         final ClusteredJob clusteredJob = ImmutableClusteredJob.builder()
                 .jobId(jobId)
                 .jobRunnerKey(jobConfig.getJobRunnerKey())
                 .schedule(quantize(jobConfig.getSchedule()))
                 .nextRunTime(nextRunTime)
                 .parameters(getParameterMapSerializer().serializeParameters(jobConfig.getParameters()))
                 .build();
 
         localJobs.remove(jobId);
         createOrReplaceWithRetry(clusteredJob);
         enqueueJob(jobId, nextRunTime);
     }
 
     private void createOrReplaceWithRetry(final ClusteredJob clusteredJob)
             throws SchedulerServiceException {
         for (int attempt = 1; attempt <= MAX_TRIES; ++attempt) {
             clusteredJobDao.delete(clusteredJob.getJobId());
             if (clusteredJobDao.create(clusteredJob)) {
                 return;
             }
         }
 
         throw new SchedulerServiceException("Unable to either create or replace clustered job: " + clusteredJob);
     }
 
 
     @Override
     public void unscheduleJob(final JobId jobId) {
         // Deliberately avoiding a short-circuit; we want to remove from both
         boolean found = localJobs.remove(jobId) != null;
         found |= clusteredJobDao.delete(jobId);
 
         // Also remove it from the queue (regardless of whether or not we think it's actually in there)
         queue.remove(jobId);
 
         if (found) {
             LOG.debug("unscheduleJob: {}", jobId);
         } else {
             LOG.debug("unscheduleJob for non-existent jobId: {}", jobId);
         }
     }
 
     @Nullable
     @Override
     public JobDetails getJobDetails(final JobId jobId) {
         final JobDetails localJob = localJobs.get(jobId);
         if (localJob != null) {
             return localJob;
         }
 
         final ClusteredJob clusteredJob = clusteredJobDao.find(jobId);
         return (clusteredJob != null) ? toJobDetails(clusteredJob) : null;
     }
 
     @Nonnull
     @Override
     public Set<JobRunnerKey> getJobRunnerKeysForAllScheduledJobs() {
         final ImmutableSet.Builder<JobRunnerKey> keys = ImmutableSet.builder();
         for (JobDetails localJob : localJobs.values()) {
             keys.add(localJob.getJobRunnerKey());
         }
         keys.addAll(clusteredJobDao.findAllJobRunnerKeys());
         return keys.build();
     }
 
     @Nonnull
     @Override
     public List<JobDetails> getJobsByJobRunnerKey(final JobRunnerKey jobRunnerKey) {
         final Map<JobId, JobDetails> jobs = new TreeMap<JobId, JobDetails>();
         for (JobDetails jobDetails : localJobs.values()) {
             if (jobDetails.getJobRunnerKey().equals(jobRunnerKey)) {
                 jobs.put(jobDetails.getJobId(), jobDetails);
             }
         }
 
         final Collection<ClusteredJob> clusteredJobs = clusteredJobDao.findByJobRunnerKey(jobRunnerKey);
         for (ClusteredJob clusteredJob : clusteredJobs) {
             jobs.put(clusteredJob.getJobId(), toJobDetails(clusteredJob));
         }
 
         return ImmutableList.copyOf(jobs.values());
     }
 
     @Nullable
     @Override
     public Date calculateNextRunTime(Schedule schedule) throws SchedulerServiceException {
         return runTimeCalculator.nextRunTime(schedule, null);
     }
 
     @Override
     protected void startImpl() throws SchedulerServiceException {
         queue.resume();
 
         if (started.compareAndSet(false, true)) {
             startWorkers();
             refreshClusteredJobs();
             scheduleRefreshJob();
         }
     }
 
 
     private void startWorkers() {
         final SchedulerQueueWorker worker = new SchedulerQueueWorker(queue, new Sink<QueuedJob>() {
             @Override
             public void consume(QueuedJob job) {
                 executeQueuedJob(job);
             }
         });
 
         int workerCount = config.workerThreadCount();
         if (workerCount <= 0) {
             workerCount = DEFAULT_WORKER_COUNT;
         }
 
         final ThreadFactory threadFactory = new WorkerThreadFactory();
         for (int i = 1; i <= workerCount; ++i) {
             threadFactory.newThread(worker).start();
         }
     }
 
     @Override
     protected void standbyImpl() throws SchedulerServiceException {
         queue.pause();
     }
 
     @Override
     protected void shutdownImpl() {
         queue.close();
     }
 
     /**
      * Preform a one-off refresh of a cluster job that is suspected of changing state.
      * <p>
      * This may be useful to an implementation that is using external messages to communicate
      * changes to clustered job schedules.
      * </p>
      *
      * @param jobId the jobId of the job whose persisted state may have changed
      */
     public void refreshClusteredJob(JobId jobId) {
         final Date nextRunTime = clusteredJobDao.getNextRunTime(jobId);
         if (nextRunTime == null) {
             if (localJobs.containsKey(jobId)) {
                 // Well, we'd better not remove it, then...
                 LOG.debug("Asked to refresh job '{}', but it is a local job so that was a bit silly.", jobId);
             } else {
                 queue.remove(jobId);
             }
             return;
         }
 
         localJobs.remove(jobId);
         try {
             queue.add(new QueuedJob(jobId, nextRunTime.getTime()));
         } catch (SchedulerQueueImpl.SchedulerShutdownException sse) {
             LOG.debug("Refresh failed for job '{}' due to scheduler shutdown", jobId, sse);
         }
     }
 
     /**
      * Perform the routine maintenance that is necessary to guarantee that this cluster node has
      * up-to-date information concerning all clustered jobs.
      * <p>
      * If the scheduler has already been shutdown, then this request is ignored.
      * </p>
      *
      * @see CaesiumSchedulerConfiguration#refreshClusteredJobsIntervalInMinutes()
      */
     public void refreshClusteredJobs() {
         final Map<JobId, Date> clusteredJobs = queue.refreshClusteredJobs();
         localJobs.keySet().removeAll(clusteredJobs.keySet());
     }
 
     private void scheduleRefreshJob() throws SchedulerServiceException {
         final int refreshInterval = config.refreshClusteredJobsIntervalInMinutes();
         if (refreshInterval > 0) {
             registerJobRunner(REFRESH_JOB_RUNNER_KEY, refreshJob);
 
             final long millis = TimeUnit.MINUTES.toMillis(refreshInterval);
             final Schedule schedule = Schedule.forInterval(millis, new Date(now() + millis));
             scheduleJob(REFRESH_JOB_ID, JobConfig.forJobRunnerKey(REFRESH_JOB_RUNNER_KEY)
                     .withRunMode(RUN_LOCALLY)
                     .withSchedule(schedule));
         } else {
             unscheduleJob(REFRESH_JOB_ID);
             unregisterJobRunner(REFRESH_JOB_RUNNER_KEY);
         }
     }
 
     /**
      * Callback from a queue worker to handle a {@code QueuedJob} that is now due.
      *
      * @param job the job to be executed
      */
     void executeQueuedJob(final QueuedJob job) {
         final JobDetails jobDetails = localJobs.get(job.getJobId());
         if (jobDetails != null) {
             executeLocalJob(jobDetails);
         } else {
             executeClusteredJob(job);
         }
     }
 
     void executeLocalJob(final JobDetails jobDetails) {
         final Date firedAt = new Date(now());
         final JobId jobId = jobDetails.getJobId();
 
         final Date scheduledRunTime = jobDetails.getNextRunTime();
         if (scheduledRunTime == null || firedAt.getTime() < scheduledRunTime.getTime()) {
             // It's very strange that a local job would launch early.  Best guess is that somebody
             // rescheduled it right in the middle of the launching process or that the clock moved.
             // Either one is pretty weird.
             LOG.debug("Launch for job '{}' either too early or after it's been deleted; scheduledRunTime={}",
                     jobDetails, scheduledRunTime);
             enqueueJob(jobId, scheduledRunTime);
             return;
         }
 
         enqueueJob(jobId, calculateNextRunTime(jobDetails, firedAt));
         launchJob(RUN_LOCALLY, firedAt, jobDetails);
     }
 
     void executeClusteredJob(final QueuedJob queuedJob) {
         final Date firedAt = new Date(now());
         final JobId jobId = queuedJob.getJobId();
         final ClusteredJob clusteredJob = clusteredJobDao.find(jobId);
         if (clusteredJob == null) {
             LOG.debug("Failed to claim '{}' for run at {}; the job no longer exists.", jobId, firedAt);
             return;
         }
 
         final JobDetails jobDetails = toJobDetails(clusteredJob);
         final Date scheduledRunTime = jobDetails.getNextRunTime();
         if (scheduledRunTime == null || queuedJob.getDeadline() < scheduledRunTime.getTime()) {
             enqueueJob(jobId, scheduledRunTime);
             return;
         }
 
         final Date nextRunTime = calculateNextRunTime(jobDetails, firedAt);
         if (!clusteredJobDao.updateNextRunTime(jobId, nextRunTime, clusteredJob.getVersion())) {
             LOG.debug("Failed to claim '{}' for run at {}; guess another node got there first?", jobId, nextRunTime);
             refreshClusteredJob(jobId);
             return;
         }
 
         enqueueJob(jobId, nextRunTime);
         launchJob(RUN_ONCE_PER_CLUSTER, firedAt, jobDetails);
     }
 
     private void launchJob(final RunMode runMode, final Date firedAt, final JobDetails jobDetails) {
         final JobLauncher launcher = new JobLauncher(this, runMode, firedAt, jobDetails.getJobId(), jobDetails);
         launcher.launch();
     }
 
     // Log errors for bad cron expressions instead of letting them kill us.  This really
     // shouldn't happen because the bad expression should have prevented the job from being
     // scheduled in the first place, but just in case...
     @Nullable
     private Date calculateNextRunTime(final JobDetails jobDetails, final Date prevRunTime) {
         try {
             return runTimeCalculator.nextRunTime(jobDetails.getSchedule(), prevRunTime);
         } catch (CronSyntaxException cse) {
             LOG.error("Clustered job '{}' has invalid cron schedule '{}' and will never run.",
                     jobDetails.getJobId(), jobDetails.getSchedule().getCronScheduleInfo().getCronExpression());
             return null;
         }
     }
 
 
     protected void enqueueJob(final JobId jobId, @Nullable final Date expectedTime) {
         try {
             if (expectedTime == null) {
                 queue.remove(jobId);
                 LOG.debug("Job '{}' has a null nextRunTime, which means we never expect it to run again", jobId);
                 return;
             }
 
             queue.add(new QueuedJob(jobId, expectedTime.getTime()));
             LOG.debug("Enqueued job '{}' for {}", jobId, expectedTime);
         } catch (SchedulerQueue.SchedulerShutdownException sse) {
             LOG.debug("Could not enqueue job '{}' because we're in the middle of shutting down", jobId, sse);
         }
     }
 
 
     @Nonnull
     JobDetails toJobDetails(final ClusteredJob clusteredJob) {
         return new LazyJobDetails(this, clusteredJob.getJobId(), clusteredJob.getJobRunnerKey(),
                 RUN_ONCE_PER_CLUSTER, clusteredJob.getSchedule(), clusteredJob.getNextRunTime(),
                 clusteredJob.getRawParameters());
     }
 
     private Schedule quantize(final Schedule schedule) {
         if (config.useFineGrainedSchedules()) {
             return schedule;
         }
 
         switch (schedule.getType()) {
             case INTERVAL: {
                 final IntervalScheduleInfo info = schedule.getIntervalScheduleInfo();
                 return Schedule.forInterval(quantizeToMinutes(info.getIntervalInMillis()), info.getFirstRunTime());
             }
 
             case CRON_EXPRESSION: {
                 final CronScheduleInfo info = schedule.getCronScheduleInfo();
                 return Schedule.forCronExpression(quantizeSecondsField(info.getCronExpression()), info.getTimeZone());
             }
 
             default:
                 throw new IllegalStateException("Unsupported schedule type: " + schedule.getType());
         }
     }
 
     /**
      * Returns {@code System.currentTimeMillis()}.
      *
      * @return {@code System.currentTimeMillis()}
      */
     @VisibleForTesting
     long now() {
         return System.currentTimeMillis();
     }
 
     /**
      * Returns the jobs that are currently in the scheduler's "to-do" list, in a map that is ordered by when they are scheduled
      * to run.
      * <p>
      * Note that this is not part of the {@code atlassian-scheduler} API and has been provided primarily as a debugging and
      * instrumentation tool.
      * </p>
      *
      * @return a map of the jobs that are currently scheduled to run, ordered by their scheduled run time.
      */
     @Internal
     public Map<JobId, Date> getPendingJobs() {
         final ImmutableMap.Builder<JobId, Date> jobMap = ImmutableMap.builder();
         for (QueuedJob job : queue.getPendingJobs()) {
             jobMap.put(job.getJobId(), new Date(job.getDeadline()));
         }
         return jobMap.build();
     }
 
     /**
      * The default factory method used for parameter map serializers.
      *
      * @param config the caesium scheduler configuration
      * @return the newly created parameter map serializer
      */
     protected static ParameterMapSerializer createParameterMapSerializer(final CaesiumSchedulerConfiguration config) {
         if (config.useQuartzJobDataMapMigration()) {
             return new LazyMigratingParameterMapSerializer();
         }
         return new ParameterMapSerializer();
     }
 
 
     class RefreshJob implements JobRunner {
         @Nullable
         @Override
         public JobRunnerResponse runJob(JobRunnerRequest request) {
             refreshClusteredJobs();
             return JobRunnerResponse.success();
         }
     }
 }