//
// Copyright © 2020, 2023 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "ArmNNProfilingServiceInitialiser.hpp"
#include "MockBackendId.hpp"
#include "ProfilingOptionsConverter.hpp"

#include <TestUtils.hpp>

#include <armnn/BackendId.hpp>
#include <armnn/Logging.hpp>

#include <armnn/profiling/ArmNNProfiling.hpp>

#include <armnn/utility/IgnoreUnused.hpp>

#include <armnnTestUtils/MockBackend.hpp>

#include <client/include/CounterIdMap.hpp>
#include <client/include/Holder.hpp>
#include <client/include/ISendTimelinePacket.hpp>
#include <client/include/ProfilingOptions.hpp>

#include <client/src/PeriodicCounterCapture.hpp>
#include <client/src/PeriodicCounterSelectionCommandHandler.hpp>
#include <client/src/ProfilingStateMachine.hpp>
#include <client/src/ProfilingUtils.hpp>
#include <client/src/RequestCounterDirectoryCommandHandler.hpp>

#include <client/src/backends/BackendProfiling.hpp>

#include <common/include/CounterDirectory.hpp>
#include <common/include/PacketVersionResolver.hpp>

#include <doctest/doctest.h>

#include <vector>
#include <cstdint>
#include <limits>

namespace arm
{

namespace pipe
{

struct LogLevelSwapper
{
public:
    LogLevelSwapper(arm::pipe::LogSeverity severity)
    {
        // Set the new log level
        arm::pipe::ConfigureLogging(true, true, severity);
    }
    ~LogLevelSwapper()
    {
        // The default log level for unit tests is "Fatal"
        arm::pipe::ConfigureLogging(true, true, arm::pipe::LogSeverity::Fatal);
    }
};

} // namespace pipe

} // namespace arm

using namespace arm::pipe;

class ReadCounterVals : public IReadCounterValues
{
    virtual bool IsCounterRegistered(uint16_t counterUid) const override
    {
        return (counterUid > 4 && counterUid < 11);
    }
    virtual bool IsCounterRegistered(const std::string& counterName) const override
    {
        armnn::IgnoreUnused(counterName);
        return false;
    }
    virtual uint16_t GetCounterCount() const override
    {
        return 1;
    }
    virtual uint32_t GetAbsoluteCounterValue(uint16_t counterUid) const override
    {
        return counterUid;
    }
    virtual uint32_t GetDeltaCounterValue(uint16_t counterUid) override
    {
        return counterUid;
    }
};

class MockBackendSendCounterPacket : public ISendCounterPacket
{
public:
    using IndexValuePairsVector = std::vector<CounterValue>;

    /// Create and write a StreamMetaDataPacket in the buffer
    virtual void SendStreamMetaDataPacket() {}

    /// Create and write a CounterDirectoryPacket from the parameters to the buffer.
    virtual void SendCounterDirectoryPacket(const ICounterDirectory& counterDirectory)
    {
        armnn::IgnoreUnused(counterDirectory);
    }

    /// Create and write a PeriodicCounterCapturePacket from the parameters to the buffer.
    virtual void SendPeriodicCounterCapturePacket(uint64_t timestamp, const IndexValuePairsVector& values)
    {
        m_timestamps.emplace_back(Timestamp{timestamp, values});
    }

    /// Create and write a PeriodicCounterSelectionPacket from the parameters to the buffer.
    virtual void SendPeriodicCounterSelectionPacket(uint32_t capturePeriod,
                                                    const std::vector<uint16_t>& selectedCounterIds)
    {
        armnn::IgnoreUnused(capturePeriod);
        armnn::IgnoreUnused(selectedCounterIds);
    }

    std::vector<Timestamp> GetTimestamps()
    {
        return  m_timestamps;
    }

    void ClearTimestamps()
    {
        m_timestamps.clear();
    }

private:
    std::vector<Timestamp> m_timestamps;
};

arm::pipe::Packet PacketWriter(uint32_t period, std::vector<uint16_t> countervalues)
{
    const uint32_t packetId = 0x40000;
    uint32_t offset = 0;
    uint32_t dataLength = static_cast<uint32_t>(4 + countervalues.size() * 2);
    std::unique_ptr<unsigned char[]> uniqueData = std::make_unique<unsigned char[]>(dataLength);
    unsigned char* data1                        = reinterpret_cast<unsigned char*>(uniqueData.get());

    WriteUint32(data1, offset, period);
    offset += 4;
    for (auto countervalue : countervalues)
    {
        WriteUint16(data1, offset, countervalue);
        offset += 2;
    }

    return {packetId, dataLength, uniqueData};
}

TEST_SUITE("BackendProfilingTestSuite")
{
TEST_CASE("BackendProfilingCounterRegisterMockBackendTest")
{
    arm::pipe::LogLevelSwapper logLevelSwapper(arm::pipe::LogSeverity::Fatal);

    // Reset the profiling service to the uninitialized state
    armnn::IRuntime::CreationOptions options;
    options.m_ProfilingOptions.m_EnableProfiling = true;

    armnn::MockBackendInitialiser initialiser;
    // Create a runtime
    armnn::RuntimeImpl runtime(options);

    // Check if the MockBackends 3 dummy counters {0, 1, 2-5 (four cores)} are registered
    armnn::BackendId mockId = armnn::MockBackendId();

    CHECK(GetProfilingService(&runtime).IsCategoryRegistered("MockCounters"));
    CHECK(GetProfilingService(&runtime).IsCounterRegistered("Mock Counter One"));
    CHECK(GetProfilingService(&runtime).IsCounterRegistered("Mock Counter Two"));
    CHECK(GetProfilingService(&runtime).IsCounterRegistered("Mock MultiCore Counter"));

    options.m_ProfilingOptions.m_EnableProfiling = false;
    GetProfilingService(&runtime).ResetExternalProfilingOptions(
        ConvertExternalProfilingOptions(options.m_ProfilingOptions), true);
}

TEST_CASE("TestBackendCounters")
{
    arm::pipe::LogLevelSwapper logLevelSwapper(arm::pipe::LogSeverity::Fatal);

    Holder holder;
    arm::pipe::PacketVersionResolver packetVersionResolver;
    ProfilingStateMachine stateMachine;
    ReadCounterVals readCounterVals;
    CounterIdMap counterIdMap;
    MockBackendSendCounterPacket sendCounterPacket;

    const std::string cpuAccId(GetComputeDeviceAsCString(armnn::Compute::CpuAcc));
    const std::string gpuAccId(GetComputeDeviceAsCString(armnn::Compute::GpuAcc));

    ProfilingOptions options;
    options.m_EnableProfiling = true;

    armnn::ArmNNProfilingServiceInitialiser initialiser;
    std::unique_ptr<IProfilingService> profilingService = arm::pipe::IProfilingService::CreateProfilingService(
        arm::pipe::MAX_ARMNN_COUNTER,
        initialiser,
        arm::pipe::ARMNN_SOFTWARE_INFO,
        arm::pipe::ARMNN_SOFTWARE_VERSION,
        arm::pipe::ARMNN_HARDWARE_VERSION);

    std::unique_ptr<IBackendProfiling> cpuBackendProfilingPtr =
        std::make_unique<BackendProfiling>(options, *profilingService.get(), cpuAccId);
    std::unique_ptr<IBackendProfiling> gpuBackendProfilingPtr =
        std::make_unique<BackendProfiling>(options, *profilingService.get(), gpuAccId);

    std::shared_ptr<IBackendProfilingContext> cpuProfilingContextPtr =
            std::make_shared<armnn::MockBackendProfilingContext>(cpuBackendProfilingPtr);
    std::shared_ptr<IBackendProfilingContext> gpuProfilingContextPtr =
            std::make_shared<armnn::MockBackendProfilingContext>(gpuBackendProfilingPtr);

    std::unordered_map<std::string,
            std::shared_ptr<IBackendProfilingContext>> backendProfilingContexts;

    backendProfilingContexts[cpuAccId] = cpuProfilingContextPtr;
    backendProfilingContexts[gpuAccId] = gpuProfilingContextPtr;

    uint16_t globalId = 5;

    counterIdMap.RegisterMapping(globalId++, 0, cpuAccId);
    counterIdMap.RegisterMapping(globalId++, 1, cpuAccId);
    counterIdMap.RegisterMapping(globalId++, 2, cpuAccId);

    counterIdMap.RegisterMapping(globalId++, 0, gpuAccId);
    counterIdMap.RegisterMapping(globalId++, 1, gpuAccId);
    counterIdMap.RegisterMapping(globalId++, 2, gpuAccId);

    backendProfilingContexts[cpuAccId] = cpuProfilingContextPtr;
    backendProfilingContexts[gpuAccId] = gpuProfilingContextPtr;

    PeriodicCounterCapture periodicCounterCapture(holder, sendCounterPacket, readCounterVals,
                                                  counterIdMap, backendProfilingContexts);

    uint16_t maxArmnnCounterId = 4;

    PeriodicCounterSelectionCommandHandler periodicCounterSelectionCommandHandler(0,
                                                  4,
                                                  packetVersionResolver.ResolvePacketVersion(0, 4).GetEncodedValue(),
                                                  backendProfilingContexts,
                                                  counterIdMap,
                                                  holder,
                                                  maxArmnnCounterId,
                                                  periodicCounterCapture,
                                                  readCounterVals,
                                                  sendCounterPacket,
                                                  stateMachine);

    stateMachine.TransitionToState(ProfilingState::NotConnected);
    stateMachine.TransitionToState(ProfilingState::WaitingForAck);
    stateMachine.TransitionToState(ProfilingState::Active);

    uint32_t period = 12345u;

    std::vector<uint16_t> cpuCounters{5, 6, 7};
    std::vector<uint16_t> gpuCounters{8, 9, 10};

    // Request only gpu counters
    periodicCounterSelectionCommandHandler(PacketWriter(period, gpuCounters));
    periodicCounterCapture.Stop();

    std::set<std::string> activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 1);
    CHECK((activeIds.find(gpuAccId) != activeIds.end()));

    std::vector<Timestamp> recievedTimestamp = sendCounterPacket.GetTimestamps();

    CHECK(recievedTimestamp[0].timestamp == period);
    CHECK(recievedTimestamp.size() == 1);
    CHECK(recievedTimestamp[0].counterValues.size() == gpuCounters.size());
    for (unsigned long i=0; i< gpuCounters.size(); ++i)
    {
        CHECK(recievedTimestamp[0].counterValues[i].counterId == gpuCounters[i]);
        CHECK(recievedTimestamp[0].counterValues[i].counterValue == i + 1u);
    }
    sendCounterPacket.ClearTimestamps();

    // Request only cpu counters
    periodicCounterSelectionCommandHandler(PacketWriter(period, cpuCounters));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 1);
    CHECK((activeIds.find(cpuAccId) != activeIds.end()));

    recievedTimestamp = sendCounterPacket.GetTimestamps();

    CHECK(recievedTimestamp[0].timestamp == period);
    CHECK(recievedTimestamp.size() == 1);
    CHECK(recievedTimestamp[0].counterValues.size() == cpuCounters.size());
    for (unsigned long i=0; i< cpuCounters.size(); ++i)
    {
        CHECK(recievedTimestamp[0].counterValues[i].counterId == cpuCounters[i]);
        CHECK(recievedTimestamp[0].counterValues[i].counterValue == i + 1u);
    }
    sendCounterPacket.ClearTimestamps();

    // Request combination of cpu & gpu counters with new period
    period = 12222u;
    periodicCounterSelectionCommandHandler(PacketWriter(period, {cpuCounters[0], gpuCounters[2],
                                                                 gpuCounters[1], cpuCounters[1], gpuCounters[0]}));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 2);
    CHECK((activeIds.find(cpuAccId) != activeIds.end()));
    CHECK((activeIds.find(gpuAccId) != activeIds.end()));

    recievedTimestamp = sendCounterPacket.GetTimestamps();
//
    CHECK(recievedTimestamp[0].timestamp == period);
    CHECK(recievedTimestamp[1].timestamp == period);

    CHECK(recievedTimestamp.size() == 2);
    CHECK(recievedTimestamp[0].counterValues.size() == 2);
    CHECK(recievedTimestamp[1].counterValues.size() == gpuCounters.size());

    CHECK(recievedTimestamp[0].counterValues[0].counterId == cpuCounters[0]);
    CHECK(recievedTimestamp[0].counterValues[0].counterValue == 1u);
    CHECK(recievedTimestamp[0].counterValues[1].counterId == cpuCounters[1]);
    CHECK(recievedTimestamp[0].counterValues[1].counterValue == 2u);

    for (unsigned long i=0; i< gpuCounters.size(); ++i)
    {
        CHECK(recievedTimestamp[1].counterValues[i].counterId == gpuCounters[i]);
        CHECK(recievedTimestamp[1].counterValues[i].counterValue == i + 1u);
    }

    sendCounterPacket.ClearTimestamps();

    // Request all counters
    std::vector<uint16_t> counterValues;
    counterValues.insert(counterValues.begin(), cpuCounters.begin(), cpuCounters.end());
    counterValues.insert(counterValues.begin(), gpuCounters.begin(), gpuCounters.end());

    periodicCounterSelectionCommandHandler(PacketWriter(period, counterValues));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 2);
    CHECK((activeIds.find(cpuAccId) != activeIds.end()));
    CHECK((activeIds.find(gpuAccId) != activeIds.end()));

    recievedTimestamp = sendCounterPacket.GetTimestamps();

    CHECK(recievedTimestamp[0].counterValues.size() == cpuCounters.size());
    for (unsigned long i=0; i< cpuCounters.size(); ++i)
    {
        CHECK(recievedTimestamp[0].counterValues[i].counterId == cpuCounters[i]);
        CHECK(recievedTimestamp[0].counterValues[i].counterValue == i + 1u);
    }

    CHECK(recievedTimestamp[1].counterValues.size() == gpuCounters.size());
    for (unsigned long i=0; i< gpuCounters.size(); ++i)
    {
        CHECK(recievedTimestamp[1].counterValues[i].counterId == gpuCounters[i]);
        CHECK(recievedTimestamp[1].counterValues[i].counterValue == i + 1u);
    }
    sendCounterPacket.ClearTimestamps();

    // Request random counters with duplicates and invalid counters
    counterValues = {0, 0, 200, cpuCounters[2], gpuCounters[0],3 ,30, cpuCounters[0],cpuCounters[2], gpuCounters[1], 3,
                     90, 0, 30, gpuCounters[0], gpuCounters[0]};

    periodicCounterSelectionCommandHandler(PacketWriter(period, counterValues));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 2);
    CHECK((activeIds.find(cpuAccId) != activeIds.end()));
    CHECK((activeIds.find(gpuAccId) != activeIds.end()));

    recievedTimestamp = sendCounterPacket.GetTimestamps();

    CHECK(recievedTimestamp.size() == 2);

    CHECK(recievedTimestamp[0].counterValues.size() == 2);

    CHECK(recievedTimestamp[0].counterValues[0].counterId == cpuCounters[0]);
    CHECK(recievedTimestamp[0].counterValues[0].counterValue == 1u);
    CHECK(recievedTimestamp[0].counterValues[1].counterId == cpuCounters[2]);
    CHECK(recievedTimestamp[0].counterValues[1].counterValue == 3u);

    CHECK(recievedTimestamp[1].counterValues.size() == 2);

    CHECK(recievedTimestamp[1].counterValues[0].counterId == gpuCounters[0]);
    CHECK(recievedTimestamp[1].counterValues[0].counterValue == 1u);
    CHECK(recievedTimestamp[1].counterValues[1].counterId == gpuCounters[1]);
    CHECK(recievedTimestamp[1].counterValues[1].counterValue == 2u);

    sendCounterPacket.ClearTimestamps();

    // Request no counters
    periodicCounterSelectionCommandHandler(PacketWriter(period, {}));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 0);

    recievedTimestamp = sendCounterPacket.GetTimestamps();
    CHECK(recievedTimestamp.size() == 0);

    sendCounterPacket.ClearTimestamps();

    // Request period of zero
    periodicCounterSelectionCommandHandler(PacketWriter(0, counterValues));
    periodicCounterCapture.Stop();

    activeIds = holder.GetCaptureData().GetActiveBackends();
    CHECK(activeIds.size() == 0);

    recievedTimestamp = sendCounterPacket.GetTimestamps();
    CHECK(recievedTimestamp.size() == 0);
}

TEST_CASE("TestBackendCounterLogging")
{
    std::stringstream ss;

    struct StreamRedirector
    {
    public:
        StreamRedirector(std::ostream &stream, std::streambuf *newStreamBuffer)
                : m_Stream(stream), m_BackupBuffer(m_Stream.rdbuf(newStreamBuffer))
        {}

        ~StreamRedirector()
        { m_Stream.rdbuf(m_BackupBuffer); }

    private:
        std::ostream &m_Stream;
        std::streambuf *m_BackupBuffer;
    };

    Holder holder;
    arm::pipe::PacketVersionResolver packetVersionResolver;
    ProfilingStateMachine stateMachine;
    ReadCounterVals readCounterVals;
    StreamRedirector redirect(std::cout, ss.rdbuf());
    CounterIdMap counterIdMap;
    MockBackendSendCounterPacket sendCounterPacket;

    const std::string cpuAccId(GetComputeDeviceAsCString(armnn::Compute::CpuAcc));
    const std::string gpuAccId(GetComputeDeviceAsCString(armnn::Compute::GpuAcc));

    ProfilingOptions options;
    options.m_EnableProfiling = true;

    armnn::ArmNNProfilingServiceInitialiser initialiser;
    std::unique_ptr<IProfilingService> profilingService = arm::pipe::IProfilingService::CreateProfilingService(
        arm::pipe::MAX_ARMNN_COUNTER,
        initialiser,
        arm::pipe::ARMNN_SOFTWARE_INFO,
        arm::pipe::ARMNN_SOFTWARE_VERSION,
        arm::pipe::ARMNN_HARDWARE_VERSION);

    std::unique_ptr<IBackendProfiling> cpuBackendProfilingPtr =
        std::make_unique<BackendProfiling>(options, *profilingService.get(), cpuAccId);

    std::shared_ptr<IBackendProfilingContext> cpuProfilingContextPtr =
            std::make_shared<armnn::MockBackendProfilingContext>(cpuBackendProfilingPtr);

    std::unordered_map<std::string,
            std::shared_ptr<IBackendProfilingContext>> backendProfilingContexts;

    uint16_t globalId = 5;
    counterIdMap.RegisterMapping(globalId, 0, cpuAccId);
    backendProfilingContexts[cpuAccId] = cpuProfilingContextPtr;

    PeriodicCounterCapture periodicCounterCapture(holder, sendCounterPacket, readCounterVals,
                                                  counterIdMap, backendProfilingContexts);

    uint16_t maxArmnnCounterId = 4;

    PeriodicCounterSelectionCommandHandler periodicCounterSelectionCommandHandler(0,
                                                  4,
                                                  packetVersionResolver.ResolvePacketVersion(0, 4).GetEncodedValue(),
                                                  backendProfilingContexts,
                                                  counterIdMap,
                                                  holder,
                                                  maxArmnnCounterId,
                                                  periodicCounterCapture,
                                                  readCounterVals,
                                                  sendCounterPacket,
                                                  stateMachine);

    stateMachine.TransitionToState(ProfilingState::NotConnected);
    stateMachine.TransitionToState(ProfilingState::WaitingForAck);
    stateMachine.TransitionToState(ProfilingState::Active);

    uint32_t period = 15939u;

    arm::pipe::SetAllLoggingSinks(true, false, false);
    arm::pipe::SetLogFilter(arm::pipe::LogSeverity::Warning);
    periodicCounterSelectionCommandHandler(PacketWriter(period, {5}));
    periodicCounterCapture.Stop();
    arm::pipe::SetLogFilter(arm::pipe::LogSeverity::Fatal);

    CHECK(ss.str().find("ActivateCounters example test error") != std::string::npos);
}

TEST_CASE("BackendProfilingContextGetSendTimelinePacket")
{
    arm::pipe::LogLevelSwapper logLevelSwapper(arm::pipe::LogSeverity::Fatal);

    // Reset the profiling service to the uninitialized state
    armnn::IRuntime::CreationOptions options;
    options.m_ProfilingOptions.m_EnableProfiling = true;

    armnn::ArmNNProfilingServiceInitialiser psInitialiser;
    std::unique_ptr<IProfilingService> profilingService = arm::pipe::IProfilingService::CreateProfilingService(
        arm::pipe::MAX_ARMNN_COUNTER,
        psInitialiser,
        arm::pipe::ARMNN_SOFTWARE_INFO,
        arm::pipe::ARMNN_SOFTWARE_VERSION,
        arm::pipe::ARMNN_HARDWARE_VERSION);

    profilingService->ConfigureProfilingService(
        ConvertExternalProfilingOptions(options.m_ProfilingOptions), true);

    armnn::MockBackendInitialiser initialiser;
    // Create a runtime. During this the mock backend will be registered and context returned.
    armnn::IRuntimePtr runtime(armnn::IRuntime::Create(options));
    armnn::MockBackendProfilingService mockProfilingService = armnn::MockBackendProfilingService::Instance();
    armnn::MockBackendProfilingContext* mockBackEndProfilingContext = mockProfilingService.GetContext();
    // Check that there is a valid context set.
    CHECK(mockBackEndProfilingContext);
    armnn::IBackendInternal::IBackendProfilingPtr& backendProfilingIface =
        mockBackEndProfilingContext->GetBackendProfiling();
    CHECK(backendProfilingIface);

    // Now for the meat of the test. We're just going to send a random packet and make sure there
    // are no exceptions or errors. The sending of packets is already tested in SendTimelinePacketTests.
    std::unique_ptr<ISendTimelinePacket> timelinePacket =
        backendProfilingIface->GetSendTimelinePacket();
    // Send TimelineEntityClassBinaryPacket
    const uint64_t entityBinaryPacketProfilingGuid = 123456u;
    timelinePacket->SendTimelineEntityBinaryPacket(entityBinaryPacketProfilingGuid);
    timelinePacket->Commit();

    // Reset the profiling servie after the test.
    options.m_ProfilingOptions.m_EnableProfiling = false;
    profilingService->ResetExternalProfilingOptions(
        ConvertExternalProfilingOptions(options.m_ProfilingOptions), true);
}

TEST_CASE("GetProfilingGuidGenerator")
{
    arm::pipe::LogLevelSwapper logLevelSwapper(arm::pipe::LogSeverity::Fatal);

    // Reset the profiling service to the uninitialized state
    armnn::IRuntime::CreationOptions options;
    options.m_ProfilingOptions.m_EnableProfiling = true;

    armnn::MockBackendInitialiser initialiser;
    // Create a runtime. During this the mock backend will be registered and context returned.
    armnn::IRuntimePtr runtime(armnn::IRuntime::Create(options));
    armnn::MockBackendProfilingService mockProfilingService = armnn::MockBackendProfilingService::Instance();
    armnn::MockBackendProfilingContext *mockBackEndProfilingContext = mockProfilingService.GetContext();
    // Check that there is a valid context set.
    CHECK(mockBackEndProfilingContext);
    armnn::IBackendInternal::IBackendProfilingPtr& backendProfilingIface =
        mockBackEndProfilingContext->GetBackendProfiling();
    CHECK(backendProfilingIface);

    // Get the Guid generator and check the getting two Guid's results in the second being greater than the first.
    IProfilingGuidGenerator& guidGenerator = backendProfilingIface->GetProfilingGuidGenerator();
    const ProfilingDynamicGuid& firstGuid = guidGenerator.NextGuid();
    const ProfilingDynamicGuid& secondGuid = guidGenerator.NextGuid();
    CHECK(secondGuid > firstGuid);

    // Reset the profiling servie after the test.
    options.m_ProfilingOptions.m_EnableProfiling = false;
}

}