/*
 * Copyright 2017 Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <folly/Optional.h>

#include <folly/fibers/FiberManagerInternal.h>
#include <folly/fibers/ForEach.h>

namespace folly {
namespace fibers {

template <class InputIterator>
typename std::vector<typename std::enable_if<
    !std::is_same<
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type,
        void>::value,
    typename std::pair<
        size_t,
        typename std::result_of<typename std::iterator_traits<
            InputIterator>::value_type()>::type>>::type>
collectN(InputIterator first, InputIterator last, size_t n) {
  typedef typename std::result_of<
      typename std::iterator_traits<InputIterator>::value_type()>::type Result;
  assert(n > 0);
  assert(std::distance(first, last) >= 0);
  assert(n <= static_cast<size_t>(std::distance(first, last)));

  struct Context {
    std::vector<std::pair<size_t, Result>> results;
    size_t tasksTodo;
    std::exception_ptr e;
    folly::Optional<Promise<void>> promise;

    Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
      this->results.reserve(tasksTodo_);
    }
  };
  auto context = std::make_shared<Context>(n);

  await([first, last, context](Promise<void> promise) mutable {
    context->promise = std::move(promise);
    for (size_t i = 0; first != last; ++i, ++first) {
#ifdef __clang__
#pragma clang diagnostic push // ignore generalized lambda capture warning
#pragma clang diagnostic ignored "-Wc++1y-extensions"
#endif
      addTask([ i, context, f = std::move(*first) ]() {
        try {
          auto result = f();
          if (context->tasksTodo == 0) {
            return;
          }
          context->results.emplace_back(i, std::move(result));
        } catch (...) {
          if (context->tasksTodo == 0) {
            return;
          }
          context->e = std::current_exception();
        }
        if (--context->tasksTodo == 0) {
          context->promise->setValue();
        }
      });
#ifdef __clang__
#pragma clang diagnostic pop
#endif
    }
  });

  if (context->e != std::exception_ptr()) {
    std::rethrow_exception(context->e);
  }

  return std::move(context->results);
}

template <class InputIterator>
typename std::enable_if<
    std::is_same<
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type,
        void>::value,
    std::vector<size_t>>::type
collectN(InputIterator first, InputIterator last, size_t n) {
  assert(n > 0);
  assert(std::distance(first, last) >= 0);
  assert(n <= static_cast<size_t>(std::distance(first, last)));

  struct Context {
    std::vector<size_t> taskIndices;
    std::exception_ptr e;
    size_t tasksTodo;
    folly::Optional<Promise<void>> promise;

    Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
      this->taskIndices.reserve(tasksTodo_);
    }
  };
  auto context = std::make_shared<Context>(n);

  await([first, last, context](Promise<void> promise) mutable {
    context->promise = std::move(promise);
    for (size_t i = 0; first != last; ++i, ++first) {
#ifdef __clang__
#pragma clang diagnostic push // ignore generalized lambda capture warning
#pragma clang diagnostic ignored "-Wc++1y-extensions"
#endif
      addTask([ i, context, f = std::move(*first) ]() {
        try {
          f();
          if (context->tasksTodo == 0) {
            return;
          }
          context->taskIndices.push_back(i);
        } catch (...) {
          if (context->tasksTodo == 0) {
            return;
          }
          context->e = std::current_exception();
        }
        if (--context->tasksTodo == 0) {
          context->promise->setValue();
        }
      });
#ifdef __clang__
#pragma clang diagnostic pop
#endif
    }
  });

  if (context->e != std::exception_ptr()) {
    std::rethrow_exception(context->e);
  }

  return context->taskIndices;
}

template <class InputIterator>
typename std::vector<
    typename std::enable_if<
        !std::is_same<
            typename std::result_of<typename std::iterator_traits<
                InputIterator>::value_type()>::type,
            void>::value,
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type>::
        type> inline collectAll(InputIterator first, InputIterator last) {
  typedef typename std::result_of<
      typename std::iterator_traits<InputIterator>::value_type()>::type Result;
  size_t n = std::distance(first, last);
  std::vector<Result> results;
  std::vector<size_t> order(n);
  results.reserve(n);

  forEach(first, last, [&results, &order](size_t id, Result result) {
    order[id] = results.size();
    results.emplace_back(std::move(result));
  });
  assert(results.size() == n);

  std::vector<Result> orderedResults;
  orderedResults.reserve(n);

  for (size_t i = 0; i < n; ++i) {
    orderedResults.emplace_back(std::move(results[order[i]]));
  }

  return orderedResults;
}

template <class InputIterator>
typename std::enable_if<
    std::is_same<
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type,
        void>::value,
    void>::type inline collectAll(InputIterator first, InputIterator last) {
  forEach(first, last, [](size_t /* id */) {});
}

template <class InputIterator>
typename std::enable_if<
    !std::is_same<
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type,
        void>::value,
    typename std::pair<
        size_t,
        typename std::result_of<typename std::iterator_traits<
            InputIterator>::value_type()>::type>>::
    type inline collectAny(InputIterator first, InputIterator last) {
  auto result = collectN(first, last, 1);
  assert(result.size() == 1);
  return std::move(result[0]);
}

template <class InputIterator>
typename std::enable_if<
    std::is_same<
        typename std::result_of<
            typename std::iterator_traits<InputIterator>::value_type()>::type,
        void>::value,
    size_t>::type inline collectAny(InputIterator first, InputIterator last) {
  auto result = collectN(first, last, 1);
  assert(result.size() == 1);
  return std::move(result[0]);
}
}
}
