package com.atlassian.stash.internal.concurrent;

import com.google.common.collect.Collections2;
import com.google.common.collect.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextRefreshedEvent;

import javax.annotation.Nonnull;
import java.util.List;
import java.util.concurrent.Executor;

/**
 * A wrapper around the standard Java {@code Executor} interface which captures the state of the current thread at the
 * time a job is executed and applies that state to and removes that state from the worker thread on which the job is
 * executed.
 */
public class StateTransferringExecutor implements Executor, ApplicationListener<ContextRefreshedEvent> {

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

    private final Executor delegate;

    private volatile List<StatefulService> services;

    public StateTransferringExecutor(Executor delegate) {
        this.delegate = delegate;

        services = ImmutableList.of();
    }

    @Override
    public void execute(@Nonnull Runnable runnable) {
        delegate.execute(wrap(runnable));
    }

    @Override
    public void onApplicationEvent(ContextRefreshedEvent event) {
        services = ImmutableList.copyOf(event.getApplicationContext().getBeansOfType(StatefulService.class).values());
    }

    protected void clearServices() {
        this.services = ImmutableList.of();
    }

    protected TransferableState getState() {
        return new CompositeTransferableState(Collections2.transform(services, StatefulService.TO_STATE));
    }

    protected Runnable wrap(Runnable runnable) {
        return new StateTransferringRunnable(runnable, getState());
    }

    protected static class StateTransferringRunnable implements Runnable {

        private final Runnable delegate;
        private final TransferableState state;

        public StateTransferringRunnable(Runnable delegate, TransferableState state) {
            this.delegate = delegate;
            this.state = state;
        }

        @Override
        public void run() {
            try {
                state.apply();
                try {
                    delegate.run();
                } finally {
                    state.remove();
                }
            } catch (RuntimeException | Error e) {
                log.error("Error while processing asynchronous task", e);
                throw e;
            }
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (!(o instanceof StateTransferringRunnable)) {
                return false;
            }

            StateTransferringRunnable that = (StateTransferringRunnable) o;

            return delegate.equals(that.delegate);
        }

        @Override
        public int hashCode() {
            return delegate.hashCode();
        }
    }
}
