package com.atlassian.bitbucket.mesh.execution;

import com.google.common.base.Throwables;
import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.Context;
import io.grpc.GrpcBackdoor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.Closeable;
import java.lang.reflect.Proxy;
import java.time.ZonedDateTime;
import java.time.temporal.ChronoField;
import java.util.List;
import java.util.Locale;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;

import static java.util.Objects.requireNonNull;
import static java.util.Optional.ofNullable;

/**
 * {@link ExecutionManager} implementation that uses gRPC's {@link Context} as the 'execution scope'. This allows us
 * to tie our {@link ExecutionContext} lifecycle to the gRPC request lifecycle when a gRPC request is being processed.
 * <p>
 * When {@link #withContext(ExecutionCallback)} is called without a gRPC request being processed, we control the
 * execution context lifecycle ourselves by create a new {@link Context.CancellableContext} and closing it when the
 * callback operation completes.
 */
public class GrpcExecutionManager implements ExecutionManager {

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

    //Key under which the GrpcExecutionContext is stored in the gRPC Context
    private static final Context.Key<GrpcExecutionContext> KEY_EXECUTION_CONTEXT = Context.key("execution-context");

    private final ConcurrentMap<String, ExecutionContext> activeContexts;
    private final AtomicLong concurrentCounter;
    private final AtomicLong contextCounter;
    private final String nodeId;

    public GrpcExecutionManager() {
        activeContexts = new ConcurrentHashMap<>();
        concurrentCounter = new AtomicLong();
        contextCounter = new AtomicLong();
        nodeId = generateNodeId();
    }

    @Nonnull
    @Override
    public Optional<ExecutionContext> getContext() {
        return ofNullable(KEY_EXECUTION_CONTEXT.get());
    }

    @Nonnull
    @Override
    public Optional<ExecutionContext> getContext(@Nonnull String contextId) {
        return ofNullable(activeContexts.get(requireNonNull(contextId, "contextId")));
    }

    @Nonnull
    @Override
    public ExecutionContext getContextOrThrow() {
        ExecutionContext context = KEY_EXECUTION_CONTEXT.get();
        if (context == null) {
            throw new IllegalStateException("No ExecutionContext is currently active");
        }
        return context;
    }

    @Nonnull
    @Override
    public <T> Callable<T> maybeWrap(@Nonnull Callable<T> callable) {
        requireNonNull(callable, "callable");

        ExecutionContext current = KEY_EXECUTION_CONTEXT.get();
        return current == null ? callable : current.wrap(callable);
    }

    @Nonnull
    @Override
    public Runnable maybeWrap(@Nonnull Runnable runnable) {
        requireNonNull(runnable, "runnable");

        ExecutionContext current = KEY_EXECUTION_CONTEXT.get();
        return current == null ? runnable : current.wrap(runnable);
    }

    @Nonnull
    @Override
    public <T> ExecutionLocal<T> newExecutionLocal() {
        return new SimpleExecutionLocal<>(null);
    }

    @Nonnull
    @Override
    public <T> ExecutionLocal<T> newExecutionLocal(@Nonnull Consumer<T> cleanupCallback) {
        return new SimpleExecutionLocal<>(requireNonNull(cleanupCallback, "cleanupCallback"));
    }

    @Override
    public <T> T withContext(@Nonnull ExecutionCallback<T> callback) {
        requireNonNull(callback, "callback");

        Context grpcContext = Context.current();
        if (!GrpcBackdoor.isCancellable(grpcContext)) {
            //Not currently executing in a gRPC context; create a cancellable context and close it when we're done
            try (Context.CancellableContext nestedContext = Context.current().withCancellation()) {
                return invokeCallbackInNewExecutionContext(callback, nestedContext);
            }
        }
        return invokeCallbackInNewExecutionContext(callback, grpcContext);
    }

    @Nonnull
    private static String generateNodeId() {
        UUID uuid = UUID.randomUUID();
        //Use radix 32 to map to characters 0-9a-v. This ensures that the resulting ID does not contain 'x', which
        //is used as a separator in the nextId().
        String nodeId = Long.toString(Math.abs(uuid.getMostSignificantBits()), 32);
        //Truncate the nodeId, because we want to keep our request IDs relatively short, and convert it to uppercase
        //so the 'x' separator will stand out more
        return nodeId.substring(0, 8).toUpperCase(Locale.ROOT);
    }

    private <T> T invokeCallbackInNewExecutionContext(ExecutionCallback<T> callback, Context grpcContext) {
        try {
            concurrentCounter.incrementAndGet();
            String contextId = nextId();

            GrpcExecutionContext executionContext = new GrpcExecutionContext(contextId, grpcContext);
            executionContext.addCleanupCallback(() -> {
                concurrentCounter.decrementAndGet();
                activeContexts.remove(contextId);
            });
            activeContexts.put(contextId, executionContext);

            //Close the execution context when the gRPC context closes
            grpcContext.addListener(context -> executionContext.close(), MoreExecutors.directExecutor());
            return executionContext.call(() -> callback.callWithContext(executionContext));
        } catch (Exception e) {
            Throwables.throwIfUnchecked(e);
            throw new RuntimeException(e);
        }
    }

    private String nextId() {
        long count = contextCounter.incrementAndGet();

        return nodeId +
                "x" + ZonedDateTime.now().get(ChronoField.MINUTE_OF_DAY) +
                "x" + count +
                "x" + concurrentCounter.get();
    }

    /**
     * {@code ExecutionContext} implementation that is stored in the {@link Context#current() current} gRPC Context.
     * {@link ExecutionLocal execution attributes} are stored inside the {@link GrpcExecutionContext}. In theory,
     * they could be stored directly in the gRPC Context, but encapsulating them here avoids bugs where attributes are
     * being retrieved/set without having an ExecutionContext set.
     */
    private static class GrpcExecutionContext implements ExecutionContext, Closeable {

        private final List<Runnable> callbacks;
        private final AtomicBoolean closed;
        private final Context grpcContext;
        private final String id;
        private final ConcurrentMap<SimpleExecutionLocal<?>, Object> localValues;

        GrpcExecutionContext(String id, Context grpcContext) {
            this.grpcContext = grpcContext.withValue(KEY_EXECUTION_CONTEXT, this);
            this.id = id;

            localValues = new ConcurrentHashMap<>();
            callbacks = new CopyOnWriteArrayList<>();
            closed = new AtomicBoolean(false);
        }

        @Override
        public void addCleanupCallback(@Nonnull Runnable callback) {
            requireNonNull(callback, "callback");
            if (closed.get()) {
                //Context is already closed, immediately execute the callback
                try {
                    callback.run();
                } catch (RuntimeException e) {
                    log.warn("Error while executing cleanup callback {}", callback.getClass().getName(), e);
                }
            } else {
                callbacks.add(callback);
            }
        }

        @Override
        public <V> V call(@Nonnull Callable<V> callable) throws Exception {
            return grpcContext.call(callable);
        }

        @Override
        public void close() {
            if (closed.compareAndSet(false, true)) {
                for (Runnable callback : callbacks) {
                    try {
                        callback.run();
                    } catch (RuntimeException e) {
                        log.warn("Error while executing cleanup callback {}", callback.getClass().getName(), e);
                    }
                }
                localValues.forEach(SimpleExecutionLocal::cleanup);
            }
        }

        @Nonnull
        @Override
        public String getId() {
            return id;
        }

        @Override
        public void run(@Nonnull Runnable runnable) {
            grpcContext.run(runnable);
        }

        @Nonnull
        @Override
        public <V> Callable<V> wrap(@Nonnull Callable<V> callable) {
            return grpcContext.wrap(callable);
        }

        @Nonnull
        @Override
        public Runnable wrap(@Nonnull Runnable runnable) {
            return grpcContext.wrap(runnable);
        }

        @SuppressWarnings("unchecked")
        @Nonnull
        @Override
        public <T> T wrap(@Nonnull T target, @Nonnull Class<?>[] interfaces) {
            requireNonNull(target, "target");
            requireNonNull(interfaces, "interfaces");
            
            return (T) Proxy.newProxyInstance(target.getClass().getClassLoader(), interfaces,
                    //Set up the current ExecutionContext and invoke the method on the target
                    (proxy, method, args) -> call(() -> method.invoke(target, args)));
        }
    }

    @SuppressWarnings({"resource", "unchecked"})
    private static class SimpleExecutionLocal<T> implements ExecutionLocal<T> {

        private final Consumer<T> cleanupCallback;

        private SimpleExecutionLocal(Consumer<T> cleanupCallback) {
            this.cleanupCallback = cleanupCallback;
        }

        @Override
        public void checkValid() {
            if (KEY_EXECUTION_CONTEXT.get() == null) {
                throw new IllegalStateException("No ExecutionContext is currently active");
            }
        }

        @Override
        public T compute(@Nonnull Function<? super T, ? extends T> mappingFunction) {
            requireNonNull(mappingFunction, "mappingFunction");
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context != null) {
                return (T) context.localValues.compute(this, (ignored, value) ->
                        requireNonNull(mappingFunction.apply((T) value), "ExecutionLocal does not support null values"));
            }
            return null;
        }

        @Override
        public T computeIfAbsent(@Nonnull Supplier<? extends T> missingValueSupplier) {
            requireNonNull(missingValueSupplier, "missingValueSupplier");
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context != null) {
                return (T) context.localValues.computeIfAbsent(this, ignored ->
                        requireNonNull(missingValueSupplier.get(), "ExecutionLocal does not support null values"));
            }
            return null;
        }

        @Override
        public T get() {
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context != null) {
                return (T) context.localValues.get(this);
            }
            return null;
        }

        @Override
        public boolean isActive() {
            return KEY_EXECUTION_CONTEXT.get() != null;
        }

        @Override
        public boolean ifPresent(@Nonnull Consumer<? super T> consumer) {
            requireNonNull(consumer, "consumer");
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context != null) {
                T localValue = (T) context.localValues.get(this);
                if (localValue != null) {
                    consumer.accept(localValue);
                    return true;
                }
            }
            return false;
        }

        @Override
        public boolean isPresent() {
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            return context != null && context.localValues.containsKey(this);
        }

        @Override
        public T remove() {
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context == null) {
                return null;
            }
            return (T) context.localValues.remove(this);
        }

        @Override
        public T removeIf(@Nonnull Predicate<? super T> predicate) {
            requireNonNull(predicate, "predicate");
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context == null) {
                return null;
            }

            for (int i = 0; i < 3; ++i) {
                T value = (T) context.localValues.get(this);
                if (value == null || !predicate.test(value)) {
                    break;
                }
                if (context.localValues.remove(this, value)) {
                    return value;
                }
            }
            //No value present, value did not match the predicate, or all attempts to remove the value failed
            return null;
        }

        @Override
        public boolean set(@Nonnull T value) {
            requireNonNull(value, "value");
            GrpcExecutionContext context = KEY_EXECUTION_CONTEXT.get();
            if (context != null) {
                context.localValues.put(this, value);
                return true;
            }
            return false;
        }

        void cleanup(@Nonnull Object value) {
            if (cleanupCallback != null) {
                try {
                    cleanupCallback.accept((T) value);
                } catch (RuntimeException e) {
                    log.warn("ExecutionLocal cleanup listener {} threw an exception",
                            cleanupCallback.getClass().getName(), e);
                }
            }
        }
    }
}
