#ifndef DOCTEST_MPI_REPORTER_H
#define DOCTEST_MPI_REPORTER_H

// #include <doctest/doctest.h>
#include <fstream>
#include <string>
#include "mpi.h"

#include <vector>
#include <mutex>

namespace doctest {

extern int nb_test_cases_skipped_insufficient_procs;
int mpi_comm_world_size();

namespace {

// https://stackoverflow.com/a/11826666/1583122
struct NullBuffer : std::streambuf {
    int overflow(int c) {
        return c;
    }
};
class NullStream : public std::ostream {
public:
    NullStream()
        : std::ostream(&nullBuff) {}

private:
    NullBuffer nullBuff = {};
};
static NullStream nullStream;

/* \brief Extends the ConsoleReporter of doctest
 *        Each process writes its results to its own file
 *        Intended to be used when a test assertion fails and the user wants to know exactly what happens on which
 * process
 */
struct MpiFileReporter : public ConsoleReporter {
    std::ofstream logfile_stream = {};

    MpiFileReporter(const ContextOptions &co)
        : ConsoleReporter(co, logfile_stream) {
        int rank = 0;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);

        std::string logfile_name = "doctest_" + std::to_string(rank) + ".log";

        logfile_stream = std::ofstream(logfile_name.c_str(), std::fstream::out);
    }
};

/* \brief Extends the ConsoleReporter of doctest
 *        Allows to manage the execution of tests in a parallel framework
 *        All results are collected on rank 0
 */
struct MpiConsoleReporter : public ConsoleReporter {
private:
    static std::ostream &replace_by_null_if_not_rank_0(std::ostream *os) {
        int rank = 0;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        if (rank == 0) {
            return *os;
        } else {
            return nullStream;
        }
    }
    std::vector<std::pair<std::string, int>> m_failure_str_queue = {};

public:
    MpiConsoleReporter(const ContextOptions &co)
        : ConsoleReporter(co, replace_by_null_if_not_rank_0(co.cout)) {}

    std::string file_line_to_string(const char *file, int line, const char *tail = "") {
        std::stringstream ss;
        ss << skipPathFromFilename(file) << (opt.gnu_file_line ? ":" : "(")
           << (opt.no_line_numbers ? 0 : line) // 0 or the real num depending on the option
           << (opt.gnu_file_line ? ":" : "):") << tail;
        return ss.str();
    }

    void test_run_end(const TestRunStats &p) override {
        ConsoleReporter::test_run_end(p);

        const bool anythingFailed = p.numTestCasesFailed > 0 || p.numAssertsFailed > 0;

        // -----------------------------------------------------
        // > Gather information in rank 0
        int n_rank, rank;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &n_rank);

        int g_numAsserts = 0;
        int g_numAssertsFailed = 0;
        int g_numTestCasesFailed = 0;

        MPI_Reduce(&p.numAsserts, &g_numAsserts, 1, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);
        MPI_Reduce(&p.numAssertsFailed, &g_numAssertsFailed, 1, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);
        MPI_Reduce(&p.numTestCasesFailed, &g_numTestCasesFailed, 1, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);

        std::vector<int> numAssertsFailedByRank;
        if (rank == 0) {
            numAssertsFailedByRank.resize(static_cast<std::size_t>(n_rank));
        }

        MPI_Gather(&p.numAssertsFailed, 1, MPI_INT, numAssertsFailedByRank.data(), 1, MPI_INT, 0, MPI_COMM_WORLD);

        if (rank == 0) {
            separator_to_stream();
            s << Color::Cyan << "[doctest] " << Color::None << "assertions on all processes: " << std::setw(6)
              << g_numAsserts << " | " << ((g_numAsserts == 0 || anythingFailed) ? Color::None : Color::Green)
              << std::setw(6) << (g_numAsserts - g_numAssertsFailed) << " passed" << Color::None << " | "
              << (g_numAssertsFailed > 0 ? Color::Red : Color::None) << std::setw(6) << g_numAssertsFailed << " failed"
              << Color::None << " |\n";
            if (nb_test_cases_skipped_insufficient_procs > 0) {
                s << Color::Cyan << "[doctest] " << Color::Yellow << "WARNING: Skipped ";
                if (nb_test_cases_skipped_insufficient_procs > 1) {
                    s << nb_test_cases_skipped_insufficient_procs << " tests requiring more than ";
                } else {
                    s << nb_test_cases_skipped_insufficient_procs << " test requiring more than ";
                }
                if (mpi_comm_world_size() > 1) {
                    s << mpi_comm_world_size() << " MPI processes to run\n";
                } else {
                    s << mpi_comm_world_size() << " MPI process to run\n";
                }
            }

            separator_to_stream();
            if (g_numAssertsFailed > 0) {

                s << Color::Cyan << "[doctest] " << Color::None << "fail on rank:" << std::setw(6) << "\n";
                for (std::size_t i = 0; i < numAssertsFailedByRank.size(); ++i) {
                    if (numAssertsFailedByRank[i] > 0) {
                        s << std::setw(16) << " -> On rank [" << i << "] with " << numAssertsFailedByRank[i]
                          << " test failed" << std::endl;
                    }
                }
            }
            s << Color::Cyan << "[doctest] " << Color::None
              << "Status: " << (g_numTestCasesFailed > 0 ? Color::Red : Color::Green)
              << ((g_numTestCasesFailed > 0) ? "FAILURE!" : "SUCCESS!") << Color::None << std::endl;
        }
    }

    void test_case_end(const CurrentTestCaseStats &st) override {
        if (is_mpi_test_case()) {
            // function called by every rank at the end of a test
            // if failed assertions happened, they have been sent to rank 0
            // here rank zero gathers them and prints them all

            int rank;
            MPI_Comm_rank(MPI_COMM_WORLD, &rank);

            std::vector<MPI_Request> requests;
            requests.reserve(m_failure_str_queue.size()); // avoid realloc & copy of MPI_Request
            for (const std::pair<std::string, int> &failure: m_failure_str_queue) {
                const std::string &failure_str = failure.first;
                const int failure_line = failure.second;

                int failure_msg_size = static_cast<int>(failure_str.size());

                requests.push_back(MPI_REQUEST_NULL);
                MPI_Isend(
                    failure_str.c_str(), failure_msg_size, MPI_BYTE, 0, failure_line, MPI_COMM_WORLD, &requests.back()
                ); // Tag = file line
            }

            // Compute the number of assert with fail among all procs
            const int nb_fail_asserts = static_cast<int>(m_failure_str_queue.size());
            int nb_fail_asserts_glob = 0;
            MPI_Reduce(&nb_fail_asserts, &nb_fail_asserts_glob, 1, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);

            if (rank == 0) {
                MPI_Status status;
                MPI_Status status_recv;

                using id_string = std::pair<int, std::string>;
                std::vector<id_string> msgs(static_cast<std::size_t>(nb_fail_asserts_glob));

                for (std::size_t i = 0; i < static_cast<std::size_t>(nb_fail_asserts_glob); ++i) {
                    MPI_Probe(MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

                    int count;
                    MPI_Get_count(&status, MPI_BYTE, &count);

                    std::string recv_msg(static_cast<std::size_t>(count), '\0');
                    void *recv_msg_data = const_cast<char *>(
                        recv_msg.data()
                    ); // const_cast needed. Non-const .data() exists in C++11 though...
                    MPI_Recv(
                        recv_msg_data, count, MPI_BYTE, status.MPI_SOURCE, status.MPI_TAG, MPI_COMM_WORLD, &status_recv
                    );

                    msgs[i] = {status.MPI_SOURCE, recv_msg};
                }

                std::sort(begin(msgs), end(msgs), [](const id_string &x, const id_string &y) {
                    return x.first < y.first;
                });

                // print
                if (nb_fail_asserts_glob > 0) {
                    separator_to_stream();
                    file_line_to_stream(tc->m_file.c_str(), static_cast<int>(tc->m_line), "\n");
                    if (tc->m_test_suite && tc->m_test_suite[0] != '\0')
                        s << Color::Yellow << "TEST SUITE: " << Color::None << tc->m_test_suite << "\n";
                    if (strncmp(tc->m_name, "  Scenario:", 11) != 0)
                        s << Color::Yellow << "TEST CASE:  ";
                    s << Color::None << tc->m_name << "\n\n";
                    for (const auto &msg: msgs) {
                        s << msg.second;
                    }
                    s << "\n";
                }
            }

            MPI_Waitall(static_cast<int>(requests.size()), requests.data(), MPI_STATUSES_IGNORE);
            m_failure_str_queue.clear();
        }

        ConsoleReporter::test_case_end(st);
    }

    bool is_mpi_test_case() const {
        return tc->m_description != nullptr && std::string(tc->m_description) == std::string("MPI_TEST_CASE");
    }

    void log_assert(const AssertData &rb) override {
        if (!is_mpi_test_case()) {
            ConsoleReporter::log_assert(rb);
        } else {
            int rank;
            MPI_Comm_rank(MPI_COMM_WORLD, &rank);

            if (!rb.m_failed && !opt.success)
                return;

            std::lock_guard<std::mutex> lock(mutex);

            std::stringstream failure_msg;
            failure_msg << Color::Red << "On rank [" << rank << "] : " << Color::None;
            failure_msg << file_line_to_string(rb.m_file, rb.m_line, " ");

            if ((rb.m_at & (assertType::is_throws_as | assertType::is_throws_with)) == 0) {
                failure_msg << Color::Cyan << assertString(rb.m_at) << "( " << rb.m_expr << " ) " << Color::None

                            << (!rb.m_failed ? "is correct!\n" : "is NOT correct!\n")
                            << "  values: " << assertString(rb.m_at) << "( " << rb.m_decomp.c_str() << " )\n";
            }

            m_failure_str_queue.push_back({failure_msg.str(), rb.m_line});
        }
    }
}; // MpiConsoleReporter

// "1" is the priority - used for ordering when multiple reporters/listeners are used
REGISTER_REPORTER("MpiConsoleReporter", 1, MpiConsoleReporter);
REGISTER_REPORTER("MpiFileReporter", 1, MpiFileReporter);

} // namespace
} // namespace doctest

#endif // DOCTEST_REPORTER_H
