/*******************************************************
 * Copyright (c) 2020, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#if !defined(AF_CPU)

#include <common/kernel_cache.hpp>

#include <common/compile_module.hpp>
#include <common/util.hpp>
#include <device_manager.hpp>
#include <platform.hpp>

#include <algorithm>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <vector>

using detail::Kernel;
using detail::Module;

using std::back_inserter;
using std::shared_timed_mutex;
using std::string;
using std::transform;
using std::unordered_map;
using std::vector;

namespace common {

using ModuleMap = unordered_map<string, Module>;

shared_timed_mutex& getCacheMutex(const int device) {
    static shared_timed_mutex mutexes[detail::DeviceManager::MAX_DEVICES];
    return mutexes[device];
}

ModuleMap& getCache(const int device) {
    static ModuleMap* caches =
        new ModuleMap[detail::DeviceManager::MAX_DEVICES];
    return caches[device];
}

Module findModule(const int device, const string& key) {
    std::shared_lock<shared_timed_mutex> readLock(getCacheMutex(device));
    auto& cache = getCache(device);
    auto iter   = cache.find(key);
    if (iter != cache.end()) { return iter->second; }
    return Module{};
}

Kernel getKernel(const string& kernelName, const vector<string>& sources,
                 const vector<TemplateArg>& targs,
                 const vector<string>& options, const bool sourceIsJIT) {
    vector<string> args;
    args.reserve(targs.size());

    transform(targs.begin(), targs.end(), back_inserter(args),
              [](const TemplateArg& arg) -> string { return arg._tparam; });

    string tInstance = kernelName;
    if (args.size() > 0) {
        tInstance = kernelName + "<" + args[0];
        for (size_t i = 1; i < args.size(); ++i) {
            tInstance += ("," + args[i]);
        }
        tInstance += ">";
    }

    const bool notJIT = !sourceIsJIT;

    vector<string> hashingVals;
    hashingVals.reserve(1 + (notJIT * (sources.size() + options.size())));
    hashingVals.push_back(tInstance);
    if (notJIT) {
        // This code path is only used for regular kernel compilation
        // since, jit funcName(kernelName) is unique to use it's hash
        // for caching the relevant compiled/linked module
        hashingVals.insert(hashingVals.end(), sources.begin(), sources.end());
        hashingVals.insert(hashingVals.end(), options.begin(), options.end());
    }

    const string moduleKey = std::to_string(deterministicHash(hashingVals));
    const int device       = detail::getActiveDeviceId();
    Module currModule      = findModule(device, moduleKey);

    if (!currModule) {
        currModule = loadModuleFromDisk(device, moduleKey, sourceIsJIT);
        if (!currModule) {
            currModule = compileModule(moduleKey, sources, options, {tInstance},
                                       sourceIsJIT);
        }

        std::unique_lock<shared_timed_mutex> writeLock(getCacheMutex(device));
        auto& cache = getCache(device);
        auto iter   = cache.find(moduleKey);
        if (iter == cache.end()) {
            // If not found, this thread is the first one to compile this
            // kernel. Keep the generated module.
            Module mod = currModule;
            getCache(device).emplace(moduleKey, mod);
        } else {
            currModule.unload();  // dump the current threads extra compilation
            currModule = iter->second;
        }
    }
#if defined(AF_CUDA)
    return getKernel(currModule, tInstance, sourceIsJIT);
#elif defined(AF_OPENCL)
    return getKernel(currModule, kernelName, sourceIsJIT);
#endif
}

}  // namespace common

#endif