// Copyright 2024 Bloomberg Finance L.P
//
// 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 <fstream>

#include <build/buildbox/local_execution.pb.h>
#include <buildboxcommon_logging.h>

#include <rumbad_statsmanager.h>
#include <rumbad_utils.h>
#include <sstream>

namespace rumbad {

StatsManager::StatsManager() : d_statsFilepath("")
{
    d_statsdata = &d_statsdata_storage;
}

void StatsManager::setStatsFilepath(const std::string &filepath)
{
    d_statsFilepath = filepath;
}

void StatsManager::incrementCompileCommandCounts(int count)
{
    d_statsdata->d_compile_command_counts += count;
    d_statsdata->updated_at = std::chrono::system_clock::now();
}

void StatsManager::incrementCompileCommandCacheHitCounts(int count)
{
    d_statsdata->d_compile_command_cache_hit_counts += count;
    d_statsdata->updated_at = std::chrono::system_clock::now();
}

void StatsManager::incrementCompileCommandCacheSkipCounts(int count)
{
    d_statsdata->d_compile_command_cache_skip_counts += count;
    d_statsdata->updated_at = std::chrono::system_clock::now();
}

void StatsManager::incrementLinkingCommandCounts(int count)
{
    d_statsdata->d_linking_command_counts += count;
    d_statsdata->updated_at = std::chrono::system_clock::now();
}

void StatsManager::incrementLinkingCommandCacheHitCounts(int count)
{
    d_statsdata->d_linking_command_cache_hit_counts += count;
    d_statsdata->updated_at = std::chrono::system_clock::now();
}

int StatsManager::totalCommandCounts()
{
    int total = d_statsdata->d_compile_command_counts +
                d_statsdata->d_linking_command_counts;
    return total;
}

float StatsManager::compileCacheHitRate()
{
    float rate = 0;
    if (d_statsdata->d_compile_command_counts > 0) {
        rate = static_cast<float>(
                   d_statsdata->d_compile_command_cache_hit_counts) /
               static_cast<float>(d_statsdata->d_compile_command_counts);
    }
    return rate;
}

float StatsManager::linkCacheHitRate()
{
    float rate = 0;
    if (d_statsdata->d_linking_command_counts > 0) {
        rate = static_cast<float>(
                   d_statsdata->d_linking_command_cache_hit_counts) /
               static_cast<float>(d_statsdata->d_linking_command_counts);
    }
    return rate;
}

StatsData const *StatsManager::readonlyStatsdata() { return d_statsdata; }

std::map<std::string, std::string> StatsManager::statsdataMap()
{
    std::map<std::string, std::string> data_map = {
        {"updated_at", formatTimestamp(d_statsdata->updated_at)},
        {"compile_command_counts",
         std::to_string(d_statsdata->d_compile_command_counts)},
        {"compile_command_cache_hit_counts",
         std::to_string(d_statsdata->d_compile_command_cache_hit_counts)},
        {"linking_command_counts",
         std::to_string(d_statsdata->d_linking_command_counts)},
        {"linking_command_cache_hit_counts",
         std::to_string(d_statsdata->d_linking_command_cache_hit_counts)},
        {"total_command_counts", std::to_string(totalCommandCounts())},
        {"compile_command_cache_hit_rate",
         std::to_string(compileCacheHitRate())},
        {"linking_command_cache_hit_rate",
         std::to_string(linkCacheHitRate())}};
    return data_map;
}

std::string StatsManager::outputString()
{
    std::stringstream outstring;
    for (auto const &[key, val] : statsdataMap()) {
        outstring << key << "=" << val << "\n";
    }
    return outstring.str();
}

void StatsManager::updateStatsInBatch(const std::vector<Message> &batch)
{
    for (size_t i = 0; i < batch.size(); i++) {
        updateStats(batch[i]);
    }
}

void StatsManager::updateStats(const Message &message)
{
    if (message.data.has_recc_data()) {
        auto counterMetrics = message.data.recc_data().counter_metrics();
        // Collect caching stats
        for (auto iter = counterMetrics.begin(); iter != counterMetrics.end();
             ++iter) {
            std::string key = (iter->first);

            // Compilation cache hit
            if (key == COUNTER_NAME_ACTION_CACHE_HIT) {
                incrementCompileCommandCounts();
                incrementCompileCommandCacheHitCounts();
            }
            // Compilation cache miss
            if (key == COUNTER_NAME_ACTION_CACHE_MISS) {
                incrementCompileCommandCounts();
            }
            // Compilation cache skip
            if (key == COUNTER_NAME_ACTION_CACHE_SKIP) {
                incrementCompileCommandCounts();
                incrementCompileCommandCacheSkipCounts();
            }
            // Link cache hit
            if (key == COUNTER_NAME_LINK_ACTION_CACHE_HIT) {
                incrementLinkingCommandCounts();
                incrementLinkingCommandCacheHitCounts();
            }
            // Link cache miss
            if (key == COUNTER_NAME_LINK_ACTION_CACHE_MISS) {
                incrementLinkingCommandCounts();
            }
        }
    }
}

std::string StatsManager::statsFilepath() { return d_statsFilepath; }

bool StatsManager::saveFilepathIsSet() { return d_statsFilepath.size() > 0; }

int StatsManager::saveStatsToFile()
{
    return saveStatsToFile(d_statsFilepath);
}

int StatsManager::saveStatsToFile(const std::string &filepath)
{
    try {
        buildboxcommon::FileUtils::writeFileAtomically(filepath,
                                                       outputString(), 0644);
        return 0;
    }
    catch (std::exception &ex) {
        BUILDBOX_LOG_WARNING("Unable to write to file: path=" + filepath);
        return 1;
    }
}

} // namespace rumbad
