package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalContextManager;
   import com.google.common.base.Function;
   import com.google.common.collect.Collections2;
   
   import java.util.Collection;
   import java.util.List;
   import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  
  /**
   * Executor service that wraps executing callables and runnables in a wrapper that transfers the thread local state of
   * the caller to the thread of the executing task.
   *
   * @since 2.0
   */
  public class ThreadLocalDelegateExecutorService<C> implements ExecutorService
  {
      protected final ThreadLocalContextManager<C> manager;
      private final ExecutorService delegate;
  
      public ThreadLocalDelegateExecutorService(ThreadLocalContextManager<C> manager, ExecutorService delegate)
      {
          this.manager = checkNotNull(manager);
          this.delegate = checkNotNull(delegate);
      }
  
      public void shutdown()
      {
          delegate.shutdown();
      }
  
      public List<Runnable> shutdownNow()
      {
          return delegate.shutdownNow();
      }
  
      public boolean isShutdown()
      {
          return delegate.isShutdown();
      }
  
      public boolean isTerminated()
      {
          return delegate.isTerminated();
      }
  
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException
      {
          return delegate.awaitTermination(timeout, unit);
      }
  
      public <T> Future<T> submit(Callable<T> callable)
      {
          return delegate.submit(threadLocalDelegateCallable(callable));
      }
  
      public <T> Future<T> submit(Runnable runnable, T t)
      {
          return delegate.submit(threadLocalDelegateRunnable(runnable), t);
      }
  
      public Future<?> submit(Runnable runnable)
      {
          return delegate.submit(threadLocalDelegateRunnable(runnable));
      }
  
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables) throws InterruptedException
      {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables));
      }
  
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables, long timeout, TimeUnit unit) throws InterruptedException
      {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables), timeout, unit);
      }
  
      public <T> T invokeAny(Collection<? extends Callable<T>> callables) throws InterruptedException, ExecutionException
      {
          return delegate.invokeAny(threadLocalDelegateCallableCollection(callables));
      }
  
      public <T> T invokeAny(Collection<? extends Callable<T>> callables, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException
      {
          return delegate.invokeAny(threadLocalDelegateCallableCollection(callables), timeout, unit);
      }
  
      public void execute(Runnable runnable)
      {
          delegate.execute(threadLocalDelegateRunnable(runnable));
      }
  
     private ThreadLocalDelegateRunnable threadLocalDelegateRunnable(Runnable runnable)
     {
         return new ThreadLocalDelegateRunnable<C>(manager, runnable);
     }
 
     private <T> ThreadLocalDelegateCallable<C, T> threadLocalDelegateCallable(Callable<T> callable)
     {
         return new ThreadLocalDelegateCallable<C, T>(manager, callable);
     }
 
     private <T> Collection<ThreadLocalDelegateCallable<C, T>> threadLocalDelegateCallableCollection(Collection<? extends Callable<T>> callables)
     {
         return Collections2.transform(callables, new Function<Callable<T>, ThreadLocalDelegateCallable<C, T>>()
         {
             public ThreadLocalDelegateCallable<C, T> apply(Callable<T> callable)
             {
                 return threadLocalDelegateCallable(callable);
             }
         });
     }
 }