// Copyright 2019 The Dawn Authors
//
// 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 "tests/perf_tests/DawnPerfTest.h"

#include "common/Assert.h"
#include "dawn_platform/tracing/TraceEvent.h"
#include "tests/perf_tests/DawnPerfTestPlatform.h"
#include "utils/Timer.h"

#include <json/value.h>
#include <json/writer.h>

#include <fstream>
#include <limits>

namespace {

    DawnPerfTestEnvironment* gTestEnv = nullptr;

    void DumpTraceEventsToJSONFile(
        const std::vector<DawnPerfTestPlatform::TraceEvent>& traceEventBuffer,
        const char* traceFile) {
        std::ofstream outFile;
        outFile.open(traceFile, std::ios_base::app);

        Json::StreamWriterBuilder builder;
        std::unique_ptr<Json::StreamWriter> writer(builder.newStreamWriter());

        for (const DawnPerfTestPlatform::TraceEvent& traceEvent : traceEventBuffer) {
            Json::Value value(Json::objectValue);

            const Json::LargestInt microseconds =
                static_cast<Json::LargestInt>(traceEvent.timestamp * 1000.0 * 1000.0);

            char phase[2] = {traceEvent.phase, '\0'};

            value["name"] = traceEvent.name;
            switch (traceEvent.category) {
                case dawn_platform::TraceCategory::General:
                    value["cat"] = "general";
                    break;
                case dawn_platform::TraceCategory::Validation:
                    value["cat"] = "validation";
                    break;
                case dawn_platform::TraceCategory::Recording:
                    value["cat"] = "recording";
                    break;
                case dawn_platform::TraceCategory::GPUWork:
                    value["cat"] = "gpu";
                    break;
                default:
                    UNREACHABLE();
            }
            value["ph"] = &phase[0];
            value["id"] = traceEvent.id;
            value["ts"] = microseconds;
            value["pid"] = "Dawn";

            outFile << ", ";
            writer->write(value, &outFile);
            outFile.flush();
        }

        outFile.close();
    }

}  // namespace

void InitDawnPerfTestEnvironment(int argc, char** argv) {
    gTestEnv = new DawnPerfTestEnvironment(argc, argv);
    DawnTestEnvironment::SetEnvironment(gTestEnv);
    testing::AddGlobalTestEnvironment(gTestEnv);
}

DawnPerfTestEnvironment::DawnPerfTestEnvironment(int argc, char** argv)
    : DawnTestEnvironment(argc, argv) {
    for (int i = 1; i < argc; ++i) {
        if (strcmp("--calibration", argv[i]) == 0) {
            mIsCalibrating = true;
            continue;
        }

        constexpr const char kOverrideStepsArg[] = "--override-steps=";
        if (strstr(argv[i], kOverrideStepsArg) == argv[i]) {
            const char* overrideSteps = argv[i] + strlen(kOverrideStepsArg);
            if (overrideSteps[0] != '\0') {
                mOverrideStepsToRun = strtoul(overrideSteps, nullptr, 0);
            }
            continue;
        }

        constexpr const char kTraceFileArg[] = "--trace-file=";
        if (strstr(argv[i], kTraceFileArg) == argv[i]) {
            const char* traceFile = argv[i] + strlen(kTraceFileArg);
            if (traceFile[0] != '\0') {
                mTraceFile = traceFile;
            }
            continue;
        }

        if (strcmp("-h", argv[i]) == 0 || strcmp("--help", argv[i]) == 0) {
            std::cout
                << "Additional flags:"
                << " [--calibration] [--override-steps=x] [--enable-tracing] [--trace-file=file]\n"
                << "  --calibration: Only run calibration. Calibration allows the perf test"
                   " runner script to save some time.\n"
                << " --override-steps: Set a fixed number of steps to run for each test\n"
                << " --trace-file: The file to dump trace results.\n"
                << std::endl;
            continue;
        }
    }
}

DawnPerfTestEnvironment::~DawnPerfTestEnvironment() = default;

void DawnPerfTestEnvironment::SetUp() {
    DawnTestEnvironment::SetUp();

    mPlatform = std::make_unique<DawnPerfTestPlatform>();
    mInstance->SetPlatform(mPlatform.get());

    // Begin writing the trace event array.
    if (mTraceFile != nullptr) {
        std::ofstream outFile;
        outFile.open(mTraceFile);
        outFile << "{ \"traceEvents\": [";
        outFile << "{}";  // Dummy object so trace events can always prepend a comma
        outFile.flush();
        outFile.close();
    }
}

void DawnPerfTestEnvironment::TearDown() {
    // End writing the trace event array.
    if (mTraceFile != nullptr) {
        std::vector<DawnPerfTestPlatform::TraceEvent> traceEventBuffer =
            mPlatform->AcquireTraceEventBuffer();

        // Write remaining trace events.
        DumpTraceEventsToJSONFile(traceEventBuffer, mTraceFile);

        std::ofstream outFile;
        outFile.open(mTraceFile, std::ios_base::app);
        outFile << "]}";
        outFile << std::endl;
        outFile.close();
    }

    DawnTestEnvironment::TearDown();
}

bool DawnPerfTestEnvironment::IsCalibrating() const {
    return mIsCalibrating;
}

unsigned int DawnPerfTestEnvironment::OverrideStepsToRun() const {
    return mOverrideStepsToRun;
}

const char* DawnPerfTestEnvironment::GetTraceFile() const {
    return mTraceFile;
}

DawnPerfTestBase::DawnPerfTestBase(DawnTestBase* test,
                                   unsigned int iterationsPerStep,
                                   unsigned int maxStepsInFlight)
    : mTest(test),
      mIterationsPerStep(iterationsPerStep),
      mMaxStepsInFlight(maxStepsInFlight),
      mTimer(utils::CreateTimer()) {
}

DawnPerfTestBase::~DawnPerfTestBase() = default;

void DawnPerfTestBase::AbortTest() {
    mRunning = false;
}

void DawnPerfTestBase::RunTest() {
    if (gTestEnv->OverrideStepsToRun() == 0) {
        // Run to compute the approximate number of steps to perform.
        mStepsToRun = std::numeric_limits<unsigned int>::max();

        // Do a warmup run for calibration.
        DoRunLoop(kCalibrationRunTimeSeconds);
        DoRunLoop(kCalibrationRunTimeSeconds);

        // Scale steps down according to the time that exceeded one second.
        double scale = kCalibrationRunTimeSeconds / mTimer->GetElapsedTime();
        mStepsToRun = static_cast<unsigned int>(static_cast<double>(mNumStepsPerformed) * scale);

        // Calibration allows the perf test runner script to save some time.
        if (gTestEnv->IsCalibrating()) {
            PrintResult("steps", mStepsToRun, "count", false);
            return;
        }
    } else {
        mStepsToRun = gTestEnv->OverrideStepsToRun();
    }

    // Do another warmup run. Seems to consistently improve results.
    DoRunLoop(kMaximumRunTimeSeconds);

    DawnPerfTestPlatform* platform =
        reinterpret_cast<DawnPerfTestPlatform*>(gTestEnv->GetInstance()->GetPlatform());
    const char* testName = ::testing::UnitTest::GetInstance()->current_test_info()->name();

    // Only enable trace event recording in this section.
    // We don't care about trace events during warmup and calibration.
    platform->EnableTraceEventRecording(true);
    {
        TRACE_EVENT0(platform, General, testName);
        for (unsigned int trial = 0; trial < kNumTrials; ++trial) {
            TRACE_EVENT0(platform, General, "Trial");
            DoRunLoop(kMaximumRunTimeSeconds);
            OutputResults();
        }
    }
    platform->EnableTraceEventRecording(false);
}

void DawnPerfTestBase::DoRunLoop(double maxRunTime) {
    dawn_platform::Platform* platform = gTestEnv->GetInstance()->GetPlatform();

    mNumStepsPerformed = 0;
    mRunning = true;

    wgpu::FenceDescriptor desc = {};
    uint64_t signaledFenceValue = 0;
    wgpu::Fence fence = mTest->queue.CreateFence(&desc);

    mTimer->Start();

    // This loop can be canceled by calling AbortTest().
    while (mRunning) {
        // Wait if there are too many steps in flight on the GPU.
        while (signaledFenceValue - fence.GetCompletedValue() >= mMaxStepsInFlight) {
            mTest->WaitABit();
        }
        TRACE_EVENT0(platform, General, "Step");
        Step();
        mTest->queue.Signal(fence, ++signaledFenceValue);

        if (mRunning) {
            ++mNumStepsPerformed;
            if (mTimer->GetElapsedTime() > maxRunTime) {
                mRunning = false;
            } else if (mNumStepsPerformed >= mStepsToRun) {
                mRunning = false;
            }
        }
    }

    // Wait for all GPU commands to complete.
    while (signaledFenceValue != fence.GetCompletedValue()) {
        mTest->WaitABit();
    }

    mTimer->Stop();
}

void DawnPerfTestBase::OutputResults() {
    DawnPerfTestPlatform* platform =
        reinterpret_cast<DawnPerfTestPlatform*>(gTestEnv->GetInstance()->GetPlatform());

    std::vector<DawnPerfTestPlatform::TraceEvent> traceEventBuffer =
        platform->AcquireTraceEventBuffer();

    struct EventTracker {
        double start = std::numeric_limits<double>::max();
        double end = 0;
        uint32_t count = 0;
    };

    EventTracker validationTracker = {};
    EventTracker recordingTracker = {};

    double totalValidationTime = 0;
    double totalRecordingTime = 0;

    // Note: We assume END timestamps always come after their corresponding BEGIN timestamps.
    // TODO(enga): When Dawn has multiple threads, stratify by thread id.
    for (const DawnPerfTestPlatform::TraceEvent& traceEvent : traceEventBuffer) {
        EventTracker* tracker = nullptr;
        double* totalTime = nullptr;

        switch (traceEvent.category) {
            case dawn_platform::TraceCategory::Validation:
                tracker = &validationTracker;
                totalTime = &totalValidationTime;
                break;
            case dawn_platform::TraceCategory::Recording:
                tracker = &recordingTracker;
                totalTime = &totalRecordingTime;
                break;
            default:
                break;
        }

        if (tracker == nullptr) {
            continue;
        }

        if (traceEvent.phase == TRACE_EVENT_PHASE_BEGIN) {
            tracker->start = std::min(tracker->start, traceEvent.timestamp);
            tracker->count++;
        }

        if (traceEvent.phase == TRACE_EVENT_PHASE_END) {
            tracker->end = std::max(tracker->end, traceEvent.timestamp);
            ASSERT(tracker->count > 0);
            tracker->count--;

            if (tracker->count == 0) {
                *totalTime += (tracker->end - tracker->start);
                *tracker = {};
            }
        }
    }

    PrintPerIterationResultFromSeconds("wall_time", mTimer->GetElapsedTime(), true);
    PrintPerIterationResultFromSeconds("validation_time", totalValidationTime, true);
    PrintPerIterationResultFromSeconds("recording_time", totalRecordingTime, true);

    const char* traceFile = gTestEnv->GetTraceFile();
    if (traceFile != nullptr) {
        DumpTraceEventsToJSONFile(traceEventBuffer, traceFile);
    }
}

void DawnPerfTestBase::PrintPerIterationResultFromSeconds(const std::string& trace,
                                                          double valueInSeconds,
                                                          bool important) const {
    if (valueInSeconds == 0) {
        return;
    }

    double secondsPerIteration =
        valueInSeconds / static_cast<double>(mNumStepsPerformed * mIterationsPerStep);

    // Give the result a different name to ensure separate graphs if we transition.
    if (secondsPerIteration > 1) {
        PrintResult(trace, secondsPerIteration * 1e3, "ms", important);
    } else if (secondsPerIteration > 1e-3) {
        PrintResult(trace, secondsPerIteration * 1e6, "us", important);
    } else {
        PrintResult(trace, secondsPerIteration * 1e9, "ns", important);
    }
}

void DawnPerfTestBase::PrintResult(const std::string& trace,
                                   double value,
                                   const std::string& units,
                                   bool important) const {
    const ::testing::TestInfo* const testInfo =
        ::testing::UnitTest::GetInstance()->current_test_info();

    const char* testName = testInfo->name();
    const char* testSuite = testInfo->test_suite_name();

    // The results are printed according to the format specified at
    // [chromium]//build/scripts/slave/performance_log_processor.py
    fflush(stdout);
    printf("%sRESULT %s%s: %s= %s%f%s %s\n", important ? "*" : "", testSuite, testName,
           trace.c_str(), "", value, "", units.c_str());
    fflush(stdout);
}

void DawnPerfTestBase::PrintResult(const std::string& trace,
                                   unsigned int value,
                                   const std::string& units,
                                   bool important) const {
    const ::testing::TestInfo* const testInfo =
        ::testing::UnitTest::GetInstance()->current_test_info();

    const char* testName = testInfo->name();
    const char* testSuite = testInfo->test_suite_name();

    // The results are printed according to the format specified at
    // [chromium]//build/scripts/slave/performance_log_processor.py
    fflush(stdout);
    printf("%sRESULT %s%s: %s= %s%u%s %s\n", important ? "*" : "", testName, testSuite,
           trace.c_str(), "", value, "", units.c_str());
    fflush(stdout);
}