package com.atlassian.stash.rest.scheduling;

import com.atlassian.sal.api.lifecycle.LifecycleAware;
import com.atlassian.scheduler.*;
import com.atlassian.scheduler.config.*;
import com.google.common.collect.ImmutableMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Date;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class CountingJobRunner implements JobRunner, LifecycleAware {

    private static final JobRunnerKey RUNNER_KEY = JobRunnerKey.of("SchedulerTest");
    private static final String CRON_JOB_PREFIX = RUNNER_KEY.toString() + ".cron";
    private static final String INTERVAL_JOB_PREFIX = RUNNER_KEY.toString() + ".interval";
    private static final String ONCE_JOB_PREFIX = RUNNER_KEY.toString() + ".once";

    private static final Logger log = LoggerFactory.getLogger(CountingJobRunner.class);

    private final ConcurrentMap<JobId, Counter> counts;

    private final SchedulerService schedulerService;

    public CountingJobRunner(SchedulerService schedulerService) {
        this.schedulerService = schedulerService;

        ImmutableMap.Builder<JobId, Counter> builder = ImmutableMap.builder();

        for (RunMode runMode : RunMode.values()) {
            builder = builder
                    .put(jobId(CRON_JOB_PREFIX, runMode), new Counter(2))
                    .put(jobId(INTERVAL_JOB_PREFIX, runMode), new Counter(2))
                    .put(jobId(ONCE_JOB_PREFIX, runMode), new Counter(1));
        }

        counts = new ConcurrentHashMap<>();
    }

    @Override
    public JobRunnerResponse runJob(JobRunnerRequest request) {
        Counter counter = counts.get(request.getJobId());
        if (counter == null) {
            return JobRunnerResponse.aborted("Unknown job id " + request.getJobId());
        }

        int value = counter.increment();
        log.info("Job {} has now run {} time(s)", request.getJobId(), value);

        return JobRunnerResponse.success();
    }

    @Override
    public void onStart() {
        schedulerService.registerJobRunner(RUNNER_KEY, this);

        for (RunMode runMode : RunMode.values()) {
            // Run every minute. Atlassian Scheduler prevents us from running more frequently than this
            scheduleJob(CRON_JOB_PREFIX, runMode, Schedule.forCronExpression("0 0/1 * * * ?"), 2);

            // Run once per minute. Atlassian Scheduler prevents us from running more frequently than this
            scheduleJob(INTERVAL_JOB_PREFIX, runMode, Schedule.forInterval(60000, new Date(System.currentTimeMillis() + 5000)), 2);

            scheduleJob(ONCE_JOB_PREFIX, runMode, Schedule.runOnce(new Date(System.currentTimeMillis() + 5000)), 1);
        }
    }

    @Override
    public void onStop() {
        // Introduced by SAL 3.x in the LifecycleAware interface.
    }

    private void scheduleJob(String jobPrefix, RunMode runMode, Schedule schedule, int minCount) {
        JobId jobId = jobId(jobPrefix, runMode);
        Counter newCounter = new Counter(minCount);
        if (counts.putIfAbsent(jobId, new Counter(minCount)) != null) {
            log.warn("Counter already registered for {}", jobId);
        }
        try {
            schedulerService.scheduleJob(jobId,
                    JobConfig.forJobRunnerKey(RUNNER_KEY)
                            .withSchedule(schedule)
                            .withRunMode(runMode));
        } catch (SchedulerServiceException e) {
            log.warn("Failed to register job {}", jobId, e);
            counts.remove(jobId, newCounter);
        }
    }

    public int awaitCount(String jobPrefix, RunMode runMode, long timeout, TimeUnit unit) {
        Counter counter = counts.get(jobId(jobPrefix, runMode));

        if (counter == null) {
            return -1;
        }

        return counter.await(timeout, unit);
    }

    private static JobId jobId(String prefix, RunMode runMode) {
        return JobId.of(prefix + " / " + runMode.name());
    }

    private static class Counter {

        private final CountDownLatch latch;
        private final AtomicInteger value;
        private final AtomicInteger waitingCount;

        private Counter(int minCount) {
            latch = new CountDownLatch(minCount);
            value = new AtomicInteger(0);
            waitingCount = new AtomicInteger(0);
        }

        public int increment() {
            int value = this.value.incrementAndGet();
            latch.countDown();
            log.info("count = {}, remaining = {}, waiting = {}", value, latch.getCount(), waitingCount.get());

            return value;
        }

        public int await(long timeout, TimeUnit unit) {
            try {
                waitingCount.incrementAndGet();
                latch.await(timeout, unit);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } finally {
                waitingCount.decrementAndGet();
            }

            return value.get();
        }

    }

}
