package com.atlassian.bitbucket.mesh.throttle;

import com.atlassian.bitbucket.mesh.config.ConfigurationUpdatedEvent;
import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.util.profiling.Ticker;
import com.atlassian.util.profiling.Timers;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import org.slf4j.Logger;
import org.springframework.scheduling.TaskScheduler;

import javax.annotation.Nonnull;
import java.math.BigInteger;
import java.time.Clock;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import static com.atlassian.bitbucket.mesh.throttle.ThrottlingOutputUtils.toHumanDuration;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.slf4j.LoggerFactory.getLogger;

public class SemaphoreThrottleService implements ThrottleService {

    private static final int MAX_ATTEMPTS = 3;
    //Tickets IDs are only used for logging purposes - they make it easier to match ticket acquired/released pairs
    private static final AtomicLong NEXT_TICKET_ID = new AtomicLong(0);

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

    private final Map<ThrottledResource, SemaphoreTicketBucket> buckets;
    private final Clock clock;
    private final TaskScheduler taskScheduler;
    private final ResourceThrottleStrategyProvider throttleStrategyProvider;
    private final ExecutionLocal<LocalTicket> tickets;

    public SemaphoreThrottleService(Clock clock, ExecutionManager executionManager, TaskScheduler taskScheduler,
                                    ResourceThrottleStrategyProvider throttleStrategyProvider) {
        this.clock = clock;
        this.taskScheduler = taskScheduler;
        this.throttleStrategyProvider = throttleStrategyProvider;

        buckets = new EnumMap<>(ThrottledResource.class);
        tickets = executionManager.newExecutionLocal(LocalTicket::release);
    }

    @Nonnull
    @Override
    public Ticket acquireTicket(@Nonnull ThrottledResource resource) {
        try (Ticker ignored = Timers.timer("acquire ticket").start(resource.getId())) {
            InvalidTicketException exception = null;

            //This loop is purely to handle the case where the LocalTicket from the ExecutionLocal is invalidated
            //concurrently. Other exceptional cases (NoSuchResourceException and ResourceBusyException) are not
            //handled in this loop and instead escape to the caller.
            for (int attempt = 0; attempt < MAX_ATTEMPTS; ++attempt) {
                LocalTicket ticket = tickets.compute(existing -> existing == null || existing.invalidated
                        ? new LocalTicket(getBucketOrThrow(resource, true), resource) : existing);
                if (ticket == null) {
                    throw new IllegalStateException("Cannot acquire a ticket without an active ExecutionContext");
                }

                try {
                    //Create a new TicketStub and return it. If a ticket has not already been acquired from the bucket,
                    //it will do so when the first stub is created. If the resource is exhausted, this will throw and
                    //_not_ retry to acquire the ticket.
                    return ticket.createStub();
                } catch (InvalidTicketException e) {
                    //The ticket was removed from the ExecutionLocal before the stub could be created. This is unlikely
                    //to happen, but it _can_ happen if another thread attempts to acquire a ticket for the same resource
                    //from the same ExecutionContext. The next iteration of the loop will retry acquiring a ticket.
                    log.trace("Invalid [{}] ticket detected.{}",
                            resource, attempt < MAX_ATTEMPTS - 1 ? " Retrying.." : "");
                    exception = e;
                }
            }

            throw exception;
        }
    }

    @Override
    public long getLongestQueueingTimeForCurrentTicketRequests(@Nonnull ThrottledResource resource) {
        long earliestQueuingTime = getBucketOrThrow(resource, true).getEarliestQueuingTime();
        return earliestQueuingTime == 0 ? 0 : clock.millis() - earliestQueuingTime;
    }

    /**
     * @return {@link TicketSummary ticket summaries} for all of the configured {@link SemaphoreTicketBucket buckets}
     */
    @Override
    @Nonnull
    public TicketSummary[] getSummaries() {
        return buckets.values().stream()
                .map(SemaphoreTicketBucket::summarize)
                .toArray(TicketSummary[]::new);
    }

    @Override
    @Nonnull
    public TicketSummary getSummary(@Nonnull ThrottledResource resource) {
        return getBucketOrThrow(resource, false).summarize();
    }

    @Override
    @Nonnull
    public Optional<Ticket> getTicket() {
        for (int attempt = 0; attempt < MAX_ATTEMPTS; ++attempt) {
            LocalTicket ticket = tickets.get();
            if (ticket == null) {
                return Optional.empty();
            }

            try {
                return Optional.ofNullable(ticket.getLatestStub());
            } catch (InvalidTicketException e) {
                //The ticket was removed from the ExecutionLocal before the stub could be stolen. This is unlikely
                //to happen, but it _can_ happen if another thread attempts to acquire a ticket from the same
                //ExecutionContext. The next iteration of the loop will retry stealing the ticket.
                log.trace("Invalid ticket detected while attempting to steal the ticket.{}",
                        attempt < MAX_ATTEMPTS - 1 ? " Retrying.." : "");
            }
        }

        return Optional.empty();
    }

    @Override
    public long getTimeSinceLastRejectedTicketRequest(@Nonnull ThrottledResource resource) {
        long lastRejectedTimestamp = getBucketOrThrow(resource, true).getLastRejectedTimestamp();
        return lastRejectedTimestamp == 0 ? 0 : System.currentTimeMillis() - lastRejectedTimestamp;
    }

    /**
     * Searches for entries in the application properties with the prefix
     * {@link ThrottlingConstants#RESOURCE_PROP_PREFIX} and configures a {@code Semaphore} for each with the specified
     * number of tickets and update strategy.
     */
    public void initialize() {
        throttleStrategyProvider.getAll().forEach((resource, strategy) -> {
            log.debug("Configured resource [{}]: {}", resource.getId(), strategy);
            buckets.put(resource, strategy.create());
        });
        maybeScheduleRefreshJob();
    }

    @EventListener
    public void onConfigUpdated(ConfigurationUpdatedEvent event) {
        if (event.getKeys().stream().anyMatch(key -> key.startsWith("throttle.") || "scaling.concurrency".equals(key))) {
            try {
                throttleStrategyProvider.reload();

                throttleStrategyProvider.getAll().forEach((resource, strategy) -> {
                    if (buckets.containsKey(resource)) {
                        strategy.update(buckets.get(resource));
                    }
                });
                log.info("Reloaded the throttling configuration");
            } catch (RuntimeException e) {
                log.warn("Failed to reload the throttling configuration", e);
            }
        }
    }

    /**
     * Retrieves a {@code TicketBucket} for throttling the named resource. If no bucket has been configured, a
     * {@link NoSuchResourceException} is thrown.
     * <p>
     * This method will <i>never</i> return {@code null}.
     *
     * @param resource the resource to retrieve the bucket for
     * @param logThrowing  {@code true} if an error should be logged; otherwise, {@code false} to skip logging
     * @return the bucket managing tickets for the specified resource
     * @throws NoSuchResourceException Thrown if no bucket has been configured for the requested resource.
     */
    @Nonnull
    @VisibleForTesting
    SemaphoreTicketBucket getBucketOrThrow(@Nonnull ThrottledResource resource, boolean logThrowing) {
        SemaphoreTicketBucket bucket = buckets.get(requireNonNull(resource, "resource"));
        if (bucket == null) {
            String message = String.format("Resource [%s] is not configured", resource);
            if (logThrowing) {
                log.error(message);
            }

            throw new NoSuchResourceException(message, resource.name());
        }
        return bucket;
    }

    private void maybeScheduleRefreshJob() {
        long intervalSeconds = 0;
        for (ResourceThrottleStrategy strategy : throttleStrategyProvider.getAll().values()) {
            long value = strategy.getUpdateInterval(SECONDS).orElse(0);
            if (value > 0) {
                intervalSeconds = intervalSeconds == 0 ? value :
                        BigInteger.valueOf(intervalSeconds).gcd(BigInteger.valueOf(value)).longValue();
            }
        }

        if (intervalSeconds > 0) {
            log.debug("Scheduling throttling adjustment task to run every {} seconds", intervalSeconds);
            taskScheduler.scheduleAtFixedRate(new UpdateThrottleStrategiesTask(), Duration.ofSeconds(intervalSeconds));
        }
    }

    private static class InvalidTicketException extends IllegalStateException {

        InvalidTicketException(String s) {
            super(s);
        }
    }

    /**
     * Ticket that is stored in the {@link ExecutionLocal}. Note that when a {@link #acquireTicket(ThrottledResource)}
     * is called, a {@link TicketStub} of the ticket is returned instead of the ticket itself. This ensures that clients
     * can close their ticket exactly once. Any attempt to close the ticket stub multiple times is ignored. When the
     * last ticket stub is closed, the {@link LocalTicket} is released and removed from the {@link ExecutionLocal}.
     */
    private class LocalTicket {

        private final SemaphoreTicketBucket bucket;
        private final long id;
        private final Deque<TicketStub> ticketStubs;
        private final ThrottledResource resource;

        private volatile long acquireTime;
        private volatile boolean invalidated;

        LocalTicket(SemaphoreTicketBucket bucket, ThrottledResource resource) {
            this.bucket = bucket;
            this.resource = resource;

            id = NEXT_TICKET_ID.incrementAndGet();
            ticketStubs = new LinkedList<>();
        }

        TicketStub createStub() {
            TicketStub stub = new TicketStub(this);
            synchronized (ticketStubs) {
                checkValidity();
                if (ticketStubs.isEmpty()) {
                    try {
                        if (bucket.tryAcquire()) {
                            acquireTime = System.nanoTime();
                            log.trace("Acquired [{}] ticket:{} ({})", resource, id, bucket);
                        } else {
                            log.warn("A [{}] ticket could not be acquired ({})", resource, bucket);
                            throw new ResourceBusyException(
                                    "The requested resource is busy and cannot service your request. " +
                                    "Please try again later.", resource);
                        }
                    } catch (InterruptedException e) {
                        log.info("Interrupted while acquiring a [{}] ticket", resource);

                        throw new ResourceBusyException("The request was interrupted before a ticket could be acquired.",
                                resource);
                    }
                }
                ticketStubs.add(stub);
            }
            return stub;
        }

        TicketStub getLatestStub() {
            synchronized (ticketStubs) {
                checkValidity();
                return ticketStubs.isEmpty() ? null : ticketStubs.getLast();
            }
        }

        void release() {
            synchronized (ticketStubs) {
                ticketStubs.clear();
                releaseAndInvalidateTicket();
            }
        }

        void returnStub(TicketStub stub) {
            synchronized (ticketStubs) {
                if (ticketStubs.remove(stub) && ticketStubs.isEmpty()) {
                    releaseAndInvalidateTicket();
                }
            }
        }

        private void checkValidity() {
            if (invalidated) {
                throw new InvalidTicketException("The ticket has been invalidated");
            }
        }

        private void releaseAndInvalidateTicket() {
            if (!invalidated) {
                SemaphoreTicketBucket bucket = buckets.get(resource);
                bucket.release();
                tickets.removeIf(ticket -> ticket == this);
                invalidated = true;
                if (log.isTraceEnabled()) {
                    String held = toHumanDuration(System.nanoTime() - acquireTime, TimeUnit.NANOSECONDS);
                    log.trace("Released [{}] ticket:{} after {} ({})", resource, id, held, bucket);
                }
            }
        }
    }

    private static class TicketStub implements Ticket {

        private final LocalTicket ticket;
        private final AtomicBoolean closed;

        TicketStub(LocalTicket ticket) {
            this.ticket = ticket;

            closed = new AtomicBoolean(false);
        }

        @Override
        public void close() {
            if (closed.compareAndSet(false, true)) {
                ticket.returnStub(this);
            }
        }

        @Nonnull
        @Override
        public ThrottledResource getResource() {
            return ticket.resource;
        }
    }

    /**
     * Task that runs every configured interval and adjusts the throttle for scm-hosting by adjusting up, down
     * or leaving unchanged the number of tickets
     */
    private class UpdateThrottleStrategiesTask implements Runnable {

        private final Map<ThrottledResource, Long> nextUpdateTimestamps;

        UpdateThrottleStrategiesTask() {
            nextUpdateTimestamps = new EnumMap<>(ThrottledResource.class);
        }

        @Override
        public void run() {
            long now = clock.millis();
            throttleStrategyProvider.getAll().forEach((resource, strategy) ->
                    strategy.getUpdateInterval(MILLISECONDS).ifPresent(updateInterval -> {
                        long nextRunTime = MoreObjects.firstNonNull(nextUpdateTimestamps.get(resource), now);
                        if (nextRunTime <= now) {
                            try {
                                strategy.update(getBucketOrThrow(resource, true));
                                nextUpdateTimestamps.put(resource, now + updateInterval);
                            } catch (RuntimeException e) {
                                log.error("Failed to update throttle limits for {}", resource, e);
                            }
                        }
                    }));
        }
    }
}
