/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 <deque>
#include <filesystem>
#include <fstream>
#include <random>
#include <string>
#include <variant>

#include "tensorrt_llm/common/logger.h"
#include "tensorrt_llm/executor/executor.h"
#include "tensorrt_llm/plugins/api/tllmPlugin.h"
#include <cxxopts.hpp>

namespace tlc = tensorrt_llm::common;
namespace tle = tensorrt_llm::executor;
namespace fs = std::filesystem;

struct RuntimeOptions
{
    std::string trtEnginePath;
    tle::SizeType32 numSysPrompts;

    tle::SizeType32 sysPromptTokens;
    tle::SizeType32 contextTokens;

    tle::SizeType32 maxTokensMean;
    tle::SizeType32 maxTokensStddev;

    tle::SizeType32 numRequests;

    size_t hostCacheSize;
    size_t maxTokensInPagedKvCache;
};

struct KVCacheBlock
{
    KVCacheBlock(size_t hash, int cacheLevel, int priority, std::optional<size_t> loraId = std::nullopt,
        std::shared_ptr<KVCacheBlock> prevBlock = nullptr);

    size_t hash;
    int cacheLevel;
    int priority;

    std::optional<size_t> loraId;

    std::shared_ptr<KVCacheBlock> prevBlock;
    std::unordered_map<size_t, std::shared_ptr<KVCacheBlock>> nextBlocks;
};

class RadixTree
{
public:
    explicit RadixTree(tle::Executor& executor);
    // Check the executor for new events.
    void pollEvents();

private:
    std::shared_ptr<tle::KVCacheEventManager> mCacheEventManager;
    // The root block of the radix tree
    std::shared_ptr<KVCacheBlock> root;
    // A table mapping block hashes to their pointers
    std::unordered_map<size_t, std::shared_ptr<KVCacheBlock>> blockTable;
    // Event counter
    size_t eventCounter;
};

// Utility function to parse input arguments
RuntimeOptions parseArgs(int argc, char* argv[]);

// Create a tle::Request
tle::Request makeRequest(int sysPromptTokens, int contextTokens, std::uniform_int_distribution<int> sysPromptSelector,
    std::normal_distribution<double> maxNumTokensSelector);

std::default_random_engine gen;

int main(int argc, char* argv[])
{
    // Register the TRT-LLM plugins
    initTrtLlmPlugins();

    auto runtimeOpts = parseArgs(argc, argv);

    // Create the executor for this engine
    auto executorConfig = tle::ExecutorConfig(1); // Beam width 1 is required for cache block reuse
    auto kvCacheConfig = tle::KvCacheConfig(true,
        runtimeOpts.maxTokensInPagedKvCache ? std::optional(runtimeOpts.maxTokensInPagedKvCache)
                                            : std::nullopt); // Enable cache block reuse
    kvCacheConfig.setHostCacheSize(runtimeOpts.hostCacheSize);
    kvCacheConfig.setEventBufferMaxSize(32768);
    executorConfig.setKvCacheConfig(kvCacheConfig);

    auto executor = tle::Executor(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig);

    auto radixTree = RadixTree(executor);

    auto activeRequests = runtimeOpts.numRequests;

    std::uniform_int_distribution sysPromptSelector(
        1, runtimeOpts.numSysPrompts); // Select a system prompt between 1 and `runtimeOpts.numSysPrompts`
    std::normal_distribution<double> maxNumTokensSelector(runtimeOpts.maxTokensMean, runtimeOpts.maxTokensStddev);

    // Create and enqueue the requests
    for (int i = 0; i < runtimeOpts.numRequests; i++)
    {
        std::ignore = executor.enqueueRequest(makeRequest(
            runtimeOpts.sysPromptTokens, runtimeOpts.contextTokens, sysPromptSelector, maxNumTokensSelector));
    }

    while (activeRequests > 0)
    {
        auto responses = executor.awaitResponses(std::chrono::milliseconds(20));
        for (auto const& response : responses)
        {
            if (response.getResult().isFinal)
                activeRequests--;
        }
        // Only call pollEvents once every 20ms. Events are only added to the queue once per iteration, so no need to
        // poll faster than this.
        radixTree.pollEvents();
    }

    return 0;
}

RuntimeOptions parseArgs(int argc, char* argv[])
{
    RuntimeOptions runtimeOpts;

    cxxopts::Options options(argv[0], "Example that demonstrates how to use the ExecutorKVCacheManager API");
    options.add_options()("h,help", "Print usage");
    options.add_options()("engine_dir", "Directory that store the engines.", cxxopts::value<std::string>());
    options.add_options()("num_sys_prompts", "Amount of unique simulated system prompts to use",
        cxxopts::value<int>()->default_value("10"));
    options.add_options()(
        "sys_prompt_tokens", "Size of the simulated system prompts", cxxopts::value<int>()->default_value("256"));
    options.add_options()("context_tokens", "Amount of varying context tokens coming after the system prompts",
        cxxopts::value<int>()->default_value("128"));
    options.add_options()(
        "max_tokens_mean", "Mean number of max output tokens", cxxopts::value<int>()->default_value("128"));
    options.add_options()(
        "max_tokens_stddev", "Standard deviation of max output tokens", cxxopts::value<int>()->default_value("32"));
    options.add_options()(
        "num_requests", "Amount of requests to send to the engine", cxxopts::value<int>()->default_value("100"));
    options.add_options()("host_cache_size", "Size of the KV Cache in host memory in bytes",
        cxxopts::value<size_t>()->default_value("0"));
    options.add_options()("max_tokens_in_paged_kv_cache", "Amount of tokens in the kv cache",
        cxxopts::value<size_t>()->default_value("0"));

    auto parsedOptions = options.parse(argc, argv);

    // Argument: help
    if (parsedOptions.count("help"))
    {
        TLLM_LOG_ERROR(options.help());
        exit(0);
    }

    // Argument: Engine directory
    if (!parsedOptions.count("engine_dir"))
    {
        TLLM_LOG_ERROR(options.help());
        TLLM_LOG_ERROR("Please specify engine directory.");
        exit(1);
    }
    runtimeOpts.trtEnginePath = parsedOptions["engine_dir"].as<std::string>();
    if (!fs::exists(runtimeOpts.trtEnginePath) || !fs::is_directory(runtimeOpts.trtEnginePath))
    {
        TLLM_LOG_ERROR("Engine directory doesn't exist.");
        exit(1);
    }

    runtimeOpts.numSysPrompts = parsedOptions["num_sys_prompts"].as<int>();
    runtimeOpts.sysPromptTokens = parsedOptions["sys_prompt_tokens"].as<int>();
    runtimeOpts.contextTokens = parsedOptions["context_tokens"].as<int>();
    runtimeOpts.maxTokensMean = parsedOptions["max_tokens_mean"].as<int>();
    runtimeOpts.maxTokensStddev = parsedOptions["max_tokens_stddev"].as<int>();
    runtimeOpts.numRequests = parsedOptions["num_requests"].as<int>();
    runtimeOpts.hostCacheSize = parsedOptions["host_cache_size"].as<size_t>();
    runtimeOpts.maxTokensInPagedKvCache = parsedOptions["max_tokens_in_paged_kv_cache"].as<size_t>();

    return runtimeOpts;
}

KVCacheBlock::KVCacheBlock(
    size_t hash, int cacheLevel, int priority, std::optional<size_t> loraId, std::shared_ptr<KVCacheBlock> prevBlock)
    : hash{hash}
    , cacheLevel{cacheLevel}
    , priority{priority}
    , loraId{loraId}
    , prevBlock{prevBlock}
    , nextBlocks{}
{
}

RadixTree::RadixTree(tle::Executor& executor)
    : mCacheEventManager(*executor.getKVCacheEventManager())
    , eventCounter{1}
{
    // Use id=-1 for the root block. Doesn't matter what exact id is used, just that it is unique.
    root = std::make_shared<KVCacheBlock>(-1, -1, -1);
    blockTable[-1] = root;

    // Wait for the `CREATED` event to be emitted.
    while (true)
    {
        auto events = mCacheEventManager->getLatestEvents();
        if (events.size() == 1)
        {
            auto const& eventData = std::get<tle::KVCacheCreatedData>(events.front().data);
            TLLM_LOG_INFO("Event ID %d: KV Cache Manager initialized with blocks per level of: %s",
                events.front().eventId, tlc::vec2str(eventData.numBlocksPerCacheLevel).c_str());
            break;
        }
    }
};

void RadixTree::pollEvents()
{
    auto events = mCacheEventManager->getLatestEvents(std::chrono::milliseconds(20));
    for (tle::KVCacheEvent const& event : events)
    {
        TLLM_CHECK(event.eventId == eventCounter++);
        if (std::holds_alternative<tle::KVCacheStoredData>(event.data))
        {
            // Blocks have been stored into the radix tree
            auto const& eventData = std::get<tle::KVCacheStoredData>(event.data);
            auto prevBlock = blockTable[eventData.parentHash.value_or(-1)];

            // This block should be in the tree
            TLLM_CHECK(blockTable.find(prevBlock->hash) != blockTable.end());

            for (auto& block : eventData.blocks)
            {

                TLLM_LOG_INFO("Event ID %d: Block %04x was inserted into the radix tree with parent %04x.",
                    event.eventId, block.blockHash, prevBlock->hash);

                // This block shouldn't already exist in the tree, and should have tokens associated with it
                TLLM_CHECK(blockTable.find(block.blockHash) == blockTable.end());
                TLLM_CHECK(block.tokens.size() > 0);

                auto thisBlock = std::make_shared<KVCacheBlock>(
                    block.blockHash, block.cacheLevel, block.priority, block.loraId, prevBlock);

                blockTable[block.blockHash] = thisBlock;
                // Link the parent to the new block
                prevBlock->nextBlocks[block.blockHash] = thisBlock;

                prevBlock = thisBlock;
            }
        }
        else if (std::holds_alternative<tle::KVCacheRemovedData>(event.data))
        {
            auto const& eventData = std::get<tle::KVCacheRemovedData>(event.data);

            for (auto const& hash : eventData.blockHashes)
            {

                TLLM_LOG_INFO("Event ID %d: Block %04x was removed from the radix tree.", event.eventId, hash);

                // This block should exist in the tree
                TLLM_CHECK(blockTable.find(hash) != blockTable.end());

                auto& block = blockTable[hash];

                // Check that the block has no children, and that the parent has the block listed as a child
                TLLM_CHECK(block->nextBlocks.size() == 0);
                TLLM_CHECK(block->prevBlock->nextBlocks.find(block->hash) != block->prevBlock->nextBlocks.end());

                // Remove the block from it's parent, and remove the entry in the block table
                block->prevBlock->nextBlocks.erase(block->hash);
                blockTable.erase(hash);
            }
        }
        else if (std::holds_alternative<tle::KVCacheUpdatedData>(event.data))
        {
            auto const& eventData = std::get<tle::KVCacheUpdatedData>(event.data);

            if (eventData.priority.has_value())
            {
                // The block priority was updated
                TLLM_LOG_INFO("Event ID %d: Block %04x priority was changed from %d to %d", event.eventId,
                    eventData.blockHash, eventData.priority->oldValue, eventData.priority->newValue);

                TLLM_CHECK(blockTable[eventData.blockHash]->priority == eventData.priority->oldValue);
                blockTable[eventData.blockHash]->priority = eventData.priority->newValue;
            }

            if (eventData.cacheLevel.has_value())
            {
                // The block cache level was updated
                TLLM_LOG_INFO("Event ID %d: Block %04x cache level was changed from %d to %d", event.eventId,
                    eventData.blockHash, eventData.cacheLevel->oldValue, eventData.cacheLevel->newValue);

                TLLM_CHECK(blockTable[eventData.blockHash]->cacheLevel == eventData.cacheLevel->oldValue);
                blockTable[eventData.blockHash]->cacheLevel = eventData.cacheLevel->newValue;
            }
        }
        else
        {
            TLLM_LOG_ERROR("Unsupported event type. This shouldn't happen!");
        }
    }
}

tle::Request makeRequest(int sysPromptTokens, int contextTokens, std::uniform_int_distribution<int> sysPromptSelector,
    std::normal_distribution<double> maxNumTokensSelector)
{
    int sysPromptVersion = sysPromptSelector(gen);
    tle::VecTokens inputTokens;

    // Add `sysPromptTokens` tokens. Add the version to the token ids to create a unique system prompt
    for (int i = 0; i < sysPromptTokens; i++)
    {
        inputTokens.emplace_back(sysPromptVersion + i);
    }
    // Add random context tokens
    for (int i = 0; i < contextTokens; i++)
    {
        inputTokens.emplace_back(rand() % 1000);
    }

    return tle::Request(inputTokens, maxNumTokensSelector(gen));
}