#include "profile.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <mutex>
#include <set>
#include <sstream>
#include <stdexcept>
#include <utility>

namespace {
    struct MpiContext {
        bool active{false};
        int rank{0};
        int size{1};
    };

    void check_mpi(
        const int result,
        const std::string_view operation
    ) {
        if (result == MPI_SUCCESS) {
            return;
        }

        std::array<char, MPI_MAX_ERROR_STRING> buffer{};
        int length = 0;
        MPI_Error_string(result, buffer.data(), &length);

        throw std::runtime_error(
            "MPI profiling operation '" + std::string(operation) +
            "' failed: " + std::string(buffer.data(), static_cast<std::size_t>(length))
        );
    }

    [[nodiscard]] MpiContext get_mpi_context(const MPI_Comm communicator) {
        int initialized = 0;
        check_mpi(MPI_Initialized(&initialized), "MPI_Initialized");

        if (initialized == 0) {
            return {};
        }

        int finalized = 0;
        check_mpi(MPI_Finalized(&finalized), "MPI_Finalized");

        if (finalized != 0) {
            return {};
        }

        if (communicator == MPI_COMM_NULL) {
            throw std::invalid_argument("Profiling aggregation requires a valid MPI communicator.");
        }

        MpiContext context{.active = true};
        check_mpi(MPI_Comm_rank(communicator, &context.rank), "MPI_Comm_rank");
        check_mpi(MPI_Comm_size(communicator, &context.size), "MPI_Comm_size");
        return context;
    }

    [[nodiscard]] std::string count_range(
        const std::uint64_t minimum,
        const std::uint64_t maximum
    ) {
        if (minimum == maximum) {
            return std::to_string(minimum);
        }
        return std::to_string(minimum) + "-" + std::to_string(maximum);
    }

    void write_csv_field(
        std::ostream &stream,
        const std::string_view field
    ) {
        stream << '"';
        for (const char character : field) {
            if (character == '"') {
                stream << "\"\"";
            } else {
                stream << character;
            }
        }
        stream << '"';
    }
} // namespace

namespace mean_field::profiling {
    struct Registry::Impl {
        struct Entry {
            std::string label;
            Statistics statistics;
        };

        mutable std::mutex mutex;
        std::map<std::string, std::size_t, std::less<>> indices;
        std::vector<Entry> entries;
    };

    Registry &Registry::Get() {
        static Registry registry;
        return registry;
    }

    Registry::Registry() : m_impl(std::make_unique<Impl>()) {
    }

    Registry::~Registry() = default;

    std::size_t Registry::Register(
        const std::string_view label,
        const std::uint64_t warmup_count
    ) {
        if (label.empty()) {
            throw std::invalid_argument("A profiling region label cannot be empty.");
        }
        if (label.find('\0') != std::string_view::npos) {
            throw std::invalid_argument("A profiling region label cannot contain a null byte.");
        }

        std::scoped_lock lock(m_impl->mutex);
        if (const auto iterator = m_impl->indices.find(label); iterator != m_impl->indices.end()) {
            Impl::Entry &entry             = m_impl->entries[iterator->second];
            entry.statistics.warmup_target = std::max(entry.statistics.warmup_target, warmup_count);
            return iterator->second;
        }

        const std::size_t index = m_impl->entries.size();
        Impl::Entry entry{.label = std::string(label)};
        entry.statistics.warmup_target = warmup_count;
        m_impl->entries.push_back(std::move(entry));
        m_impl->indices.emplace(m_impl->entries.back().label, index);
        return index;
    }

    void Registry::Record(
        const std::string_view label,
        const double seconds,
        const std::uint64_t warmup_count
    ) {
        if (!std::isfinite(seconds) || seconds < 0.0) {
            throw std::invalid_argument("A profiling duration must be finite and nonnegative.");
        }

        const std::size_t region = Register(label, warmup_count);
        std::scoped_lock lock(m_impl->mutex);
        Statistics &statistics = m_impl->entries[region].statistics;
        const bool is_warmup   = statistics.observations < statistics.warmup_target;
        ++statistics.observations;

        if (is_warmup) {
            ++statistics.warmups;
            return;
        }

        ++statistics.samples;
        statistics.total_seconds += seconds;
        if (statistics.samples == 1) {
            statistics.minimum_seconds = seconds;
            statistics.maximum_seconds = seconds;
        } else {
            statistics.minimum_seconds = std::min(statistics.minimum_seconds, seconds);
            statistics.maximum_seconds = std::max(statistics.maximum_seconds, seconds);
        }
    }

    void Registry::AddCount(
        const std::string_view label,
        const std::uint64_t work_units
    ) {
        const std::size_t region = Register(label, 0);
        std::scoped_lock lock(m_impl->mutex);
        Statistics &statistics = m_impl->entries[region].statistics;
        if (work_units > std::numeric_limits<std::uint64_t>::max() - statistics.work_units) {
            throw std::overflow_error("A profiling work counter overflowed.");
        }
        statistics.work_units += work_units;
    }

    void Registry::Record(
        const std::size_t region,
        const double seconds
    ) noexcept {
        if (!std::isfinite(seconds) || seconds < 0.0) {
            return;
        }

        try {
            std::scoped_lock lock(m_impl->mutex);
            if (region >= m_impl->entries.size()) {
                return;
            }

            Statistics &statistics = m_impl->entries[region].statistics;
            const bool is_warmup   = statistics.observations < statistics.warmup_target;
            ++statistics.observations;

            if (is_warmup) {
                ++statistics.warmups;
                return;
            }

            ++statistics.samples;
            statistics.total_seconds += seconds;
            if (statistics.samples == 1) {
                statistics.minimum_seconds = seconds;
                statistics.maximum_seconds = seconds;
            } else {
                statistics.minimum_seconds = std::min(statistics.minimum_seconds, seconds);
                statistics.maximum_seconds = std::max(statistics.maximum_seconds, seconds);
            }
        } catch (...) {
        }
    }

    void Registry::AddCount(
        const std::size_t region,
        const std::uint64_t work_units
    ) noexcept {
        try {
            std::scoped_lock lock(m_impl->mutex);
            if (region >= m_impl->entries.size()) {
                return;
            }

            Statistics &statistics = m_impl->entries[region].statistics;
            if (work_units > std::numeric_limits<std::uint64_t>::max() - statistics.work_units) {
                statistics.work_units = std::numeric_limits<std::uint64_t>::max();
            } else {
                statistics.work_units += work_units;
            }
        } catch (...) {
        }
    }

    void Registry::Reset() {
        std::scoped_lock lock(m_impl->mutex);
        for (Impl::Entry &entry : m_impl->entries) {
            const std::uint64_t warmup_target = entry.statistics.warmup_target;
            entry.statistics                  = {};
            entry.statistics.warmup_target    = warmup_target;
        }
    }

    std::map<
        std::string,
        Statistics,
        std::less<>>
    Registry::Snapshot() const {
        std::map<std::string, Statistics, std::less<>> snapshot;
        std::scoped_lock lock(m_impl->mutex);
        for (const Impl::Entry &entry : m_impl->entries) {
            snapshot.emplace(entry.label, entry.statistics);
        }
        return snapshot;
    }

    std::vector<DistributedStatistics> Registry::Aggregate(const MPI_Comm communicator) const {
        const std::map<std::string, Statistics, std::less<>> local_snapshot = Snapshot();
        const MpiContext mpi_context                                        = get_mpi_context(communicator);

        std::vector<std::string> labels;
        if (!mpi_context.active) {
            labels.reserve(local_snapshot.size());
            for (const auto &[label, statistics] : local_snapshot) {
                (void)statistics;
                labels.push_back(label);
            }
        } else {
            std::string serialized_labels;
            for (const auto &[label, statistics] : local_snapshot) {
                (void)statistics;
                serialized_labels.append(label);
                serialized_labels.push_back('\0');
            }

            if (serialized_labels.size() > static_cast<std::size_t>(std::numeric_limits<int>::max())) {
                throw std::overflow_error("The local profiling label table is too large for MPI_Allgatherv.");
            }

            const int local_bytes = static_cast<int>(serialized_labels.size());
            std::vector<int> byte_counts(static_cast<std::size_t>(mpi_context.size));
            check_mpi(
                MPI_Allgather(&local_bytes, 1, MPI_INT, byte_counts.data(), 1, MPI_INT, communicator),
                "MPI_Allgather(profile label sizes)"
            );

            std::vector<int> displacements(static_cast<std::size_t>(mpi_context.size));
            int total_bytes = 0;
            for (int rank = 0; rank < mpi_context.size; ++rank) {
                if (byte_counts[rank] < 0 || byte_counts[rank] > std::numeric_limits<int>::max() - total_bytes) {
                    throw std::overflow_error("The distributed profiling label table is too large for MPI_Allgatherv.");
                }
                displacements[rank] = total_bytes;
                total_bytes += byte_counts[rank];
            }

            std::vector<char> all_serialized_labels(static_cast<std::size_t>(total_bytes));
            check_mpi(
                MPI_Allgatherv(
                    serialized_labels.data(), local_bytes, MPI_CHAR, all_serialized_labels.data(), byte_counts.data(),
                    displacements.data(), MPI_CHAR, communicator
                ),
                "MPI_Allgatherv(profile labels)"
            );

            std::set<std::string, std::less<>> unique_labels;
            for (int rank = 0; rank < mpi_context.size; ++rank) {
                const char *position = all_serialized_labels.data() + displacements[rank];
                const char *end      = position + byte_counts[rank];
                while (position != end) {
                    const void *terminator_address =
                        std::memchr(position, '\0', static_cast<std::size_t>(end - position));
                    if (terminator_address == nullptr) {
                        throw std::runtime_error("A distributed profiling label table is malformed.");
                    }
                    const auto *terminator = static_cast<const char *>(terminator_address);
                    unique_labels.emplace(position, terminator);
                    position = terminator + 1;
                }
            }
            labels.assign(unique_labels.begin(), unique_labels.end());
        }

        std::vector<DistributedStatistics> aggregate(labels.size());
        if (labels.empty()) {
            return aggregate;
        }

        std::vector<std::uint64_t> local_samples(labels.size(), 0);
        std::vector<std::uint64_t> local_warmups(labels.size(), 0);
        std::vector<std::uint64_t> local_work_units(labels.size(), 0);
        std::vector<double> local_averages(labels.size(), 0.0);
        std::vector<double> local_minima(labels.size(), std::numeric_limits<double>::infinity());
        std::vector<double> local_maxima(labels.size(), 0.0);
        std::vector<double> local_totals(labels.size(), 0.0);

        for (std::size_t index = 0; index < labels.size(); ++index) {
            if (const auto iterator = local_snapshot.find(labels[index]); iterator != local_snapshot.end()) {
                const Statistics &statistics = iterator->second;
                local_samples[index]         = statistics.samples;
                local_warmups[index]         = statistics.warmups;
                local_work_units[index]      = statistics.work_units;
                local_totals[index]          = statistics.total_seconds;
                if (statistics.samples != 0) {
                    local_averages[index] = statistics.total_seconds / static_cast<double>(statistics.samples);
                    local_minima[index]   = statistics.minimum_seconds;
                    local_maxima[index]   = statistics.maximum_seconds;
                }
            }
        }

        std::vector<std::uint64_t> minimum_samples    = local_samples;
        std::vector<std::uint64_t> maximum_samples    = local_samples;
        std::vector<std::uint64_t> maximum_warmups    = local_warmups;
        std::vector<std::uint64_t> minimum_work_units = local_work_units;
        std::vector<std::uint64_t> maximum_work_units = local_work_units;
        std::vector<double> maximum_rank_averages     = local_averages;
        std::vector<double> global_minima             = local_minima;
        std::vector<double> global_maxima             = local_maxima;
        std::vector<double> maximum_rank_totals       = local_totals;

        if (mpi_context.active) {
            if (labels.size() > static_cast<std::size_t>(std::numeric_limits<int>::max())) {
                throw std::overflow_error("There are too many profiling regions for one MPI reduction.");
            }
            const int count = static_cast<int>(labels.size());

            check_mpi(
                MPI_Allreduce(local_samples.data(), minimum_samples.data(), count, MPI_UINT64_T, MPI_MIN, communicator),
                "MPI_Allreduce(minimum profile samples)"
            );
            check_mpi(
                MPI_Allreduce(local_samples.data(), maximum_samples.data(), count, MPI_UINT64_T, MPI_MAX, communicator),
                "MPI_Allreduce(maximum profile samples)"
            );
            check_mpi(
                MPI_Allreduce(local_warmups.data(), maximum_warmups.data(), count, MPI_UINT64_T, MPI_MAX, communicator),
                "MPI_Allreduce(profile warmups)"
            );
            check_mpi(
                MPI_Allreduce(
                    local_work_units.data(), minimum_work_units.data(), count, MPI_UINT64_T, MPI_MIN, communicator
                ),
                "MPI_Allreduce(minimum profile work)"
            );
            check_mpi(
                MPI_Allreduce(
                    local_work_units.data(), maximum_work_units.data(), count, MPI_UINT64_T, MPI_MAX, communicator
                ),
                "MPI_Allreduce(maximum profile work)"
            );
            check_mpi(
                MPI_Allreduce(
                    local_averages.data(), maximum_rank_averages.data(), count, MPI_DOUBLE, MPI_MAX, communicator
                ),
                "MPI_Allreduce(profile averages)"
            );
            check_mpi(
                MPI_Allreduce(local_minima.data(), global_minima.data(), count, MPI_DOUBLE, MPI_MIN, communicator),
                "MPI_Allreduce(profile minima)"
            );
            check_mpi(
                MPI_Allreduce(local_maxima.data(), global_maxima.data(), count, MPI_DOUBLE, MPI_MAX, communicator),
                "MPI_Allreduce(profile maxima)"
            );
            check_mpi(
                MPI_Allreduce(
                    local_totals.data(), maximum_rank_totals.data(), count, MPI_DOUBLE, MPI_MAX, communicator
                ),
                "MPI_Allreduce(profile totals)"
            );
        }

        for (std::size_t index = 0; index < labels.size(); ++index) {
            aggregate[index] = {
                .label                        = labels[index],
                .minimum_samples              = minimum_samples[index],
                .maximum_samples              = maximum_samples[index],
                .maximum_warmups              = maximum_warmups[index],
                .minimum_work_units           = minimum_work_units[index],
                .maximum_work_units           = maximum_work_units[index],
                .maximum_rank_average_seconds = maximum_rank_averages[index],
                .global_minimum_seconds       = std::isfinite(global_minima[index]) ? global_minima[index] : 0.0,
                .global_maximum_seconds       = global_maxima[index],
                .maximum_rank_total_seconds   = maximum_rank_totals[index]
            };
        }

        return aggregate;
    }

    void Registry::Print(
        const MPI_Comm communicator,
        std::ostream &stream
    ) const {
        const std::vector<DistributedStatistics> aggregate = Aggregate(communicator);
        const MpiContext mpi_context                       = get_mpi_context(communicator);
        if (mpi_context.rank != 0) {
            return;
        }

        std::ios old_state(nullptr);
        old_state.copyfmt(stream);

        stream << '\n';
        stream << std::left << std::setw(58) << "Profile Region" << std::right << std::setw(13) << "Samples"
               << std::setw(11) << "Warmups" << std::setw(15) << "Work/rank" << std::setw(14) << "Avg max ms"
               << std::setw(14) << "Min ms" << std::setw(14) << "Max ms" << std::setw(14) << "Total max s" << '\n';
        stream << std::string(153, '-') << '\n';

        for (const DistributedStatistics &statistics : aggregate) {
            stream << std::left << std::setw(58) << statistics.label << std::right << std::setw(13)
                   << count_range(statistics.minimum_samples, statistics.maximum_samples) << std::setw(11)
                   << statistics.maximum_warmups << std::setw(15)
                   << count_range(statistics.minimum_work_units, statistics.maximum_work_units) << std::setw(14)
                   << std::fixed << std::setprecision(3) << 1.0e3 * statistics.maximum_rank_average_seconds
                   << std::setw(14) << 1.0e3 * statistics.global_minimum_seconds << std::setw(14)
                   << 1.0e3 * statistics.global_maximum_seconds << std::setw(14) << std::setprecision(6)
                   << statistics.maximum_rank_total_seconds << '\n';
        }

        stream << std::string(153, '=') << '\n';
        stream << "MPI ranks: " << mpi_context.size << "\n\n";
        stream.copyfmt(old_state);
    }

    void Registry::Print(const MPI_Comm communicator) const {
        Print(communicator, std::cout);
    }

    void Registry::PrintCsv(
        const MPI_Comm communicator,
        std::ostream &stream
    ) const {
        const std::vector<DistributedStatistics> aggregate = Aggregate(communicator);
        const MpiContext mpi_context                       = get_mpi_context(communicator);
        if (mpi_context.rank != 0) {
            return;
        }

        stream << "label,minimum_samples,maximum_samples,maximum_warmups,minimum_work_units,maximum_work_units,"
                  "maximum_rank_average_seconds,global_minimum_seconds,global_maximum_seconds,"
                  "maximum_rank_total_seconds,mpi_ranks\n";

        for (const DistributedStatistics &statistics : aggregate) {
            write_csv_field(stream, statistics.label);
            stream << ',' << statistics.minimum_samples << ',' << statistics.maximum_samples << ','
                   << statistics.maximum_warmups << ',' << statistics.minimum_work_units << ','
                   << statistics.maximum_work_units << ',' << std::setprecision(17)
                   << statistics.maximum_rank_average_seconds << ',' << statistics.global_minimum_seconds << ','
                   << statistics.global_maximum_seconds << ',' << statistics.maximum_rank_total_seconds << ','
                   << mpi_context.size << '\n';
        }
    }

    Region::Region(
        const std::string_view label,
        const std::uint64_t warmup_count
    )
        : m_region(
              Registry::Get().Register(
                  label,
                  warmup_count
              )
          ) {
    }

    void Region::Record(const double seconds) const noexcept {
        Registry::Get().Record(m_region, seconds);
    }

    void Region::AddCount(const std::uint64_t work_units) const noexcept {
        Registry::Get().AddCount(m_region, work_units);
    }

    ScopedTimer::ScopedTimer(const Region &region) noexcept
        : m_region(region),
          m_start(std::chrono::steady_clock::now()) {
    }

    ScopedTimer::~ScopedTimer() noexcept {
        const auto stop = std::chrono::steady_clock::now();
        m_region.Record(std::chrono::duration<double>(stop - m_start).count());
    }
} // namespace mean_field::profiling