// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once

#include <carb/assets/IAssets.h>
#include <carb/datasource/DataSourceUtils.h>
#include <carb/datasource/IDataSource.h>
#include <carb/filesystem/IFileSystem.h>
#include <carb/logging/Log.h>
#include <omni/usd/PathUtils.h>

#include <unordered_map>
#include <vector>

namespace omni
{
namespace renderer
{
namespace utils
{
inline carb::datasource::IDataSource* getDataSourceFile()
{
    return carb::getFramework()->tryAcquireInterface<carb::datasource::IDataSource>("carb.datasource-file.plugin");
}

static bool isPathPrefixed(const std::string& path, const char* prefix, size_t prefixLen)
{
    return path.length() > prefixLen && strncasecmp(path.c_str(), prefix, prefixLen) == 0;
}

inline std::vector<std::tuple<char, carb::datasource::IDataSource*, carb::datasource::Connection*>> getFilesystemConnections()
{
    carb::filesystem::IFileSystem* fs = carb::getFramework()->acquireInterface<carb::filesystem::IFileSystem>();
    carb::datasource::IDataSource* dataSource =
        carb::getFramework()->tryAcquireInterface<carb::datasource::IDataSource>("carb.datasource-omniclient.plugin");

    if (!dataSource)
    {
        dataSource = getDataSourceFile();
    }

    std::vector<std::tuple<char, carb::datasource::IDataSource*, carb::datasource::Connection*>> vec;

    if (fs == nullptr || dataSource == nullptr)
        return vec;

#if CARB_PLATFORM_WINDOWS

    for (char drive = 'A'; drive <= 'Z'; drive++)
    {
        std::string drivePath = std::string(1, drive) + ":";
        if (fs->exists(drivePath.c_str()))
        {
            carb::datasource::Connection* connection =
                carb::datasource::connectAndWait(carb::datasource::ConnectionDesc{ drivePath.c_str() }, dataSource);
            if (connection == nullptr)
            {
                CARB_LOG_ERROR("failed to get connection for drive %c", drive);
                continue;
            }

            vec.push_back(std::make_tuple(drive, dataSource, connection));
        }
    }

#else
    // Just add the root of the filesystem
    carb::datasource::Connection* connection =
        carb::datasource::connectAndWait(carb::datasource::ConnectionDesc{ "/" }, dataSource);
    if (connection == nullptr)
        CARB_LOG_ERROR("failed to get connection for the filesystem root");
    else
        vec.push_back(std::make_tuple('/', dataSource, connection));
#endif

    return vec;
}

class AssetsManager
{
public:
    struct DataSourceMapping
    {
        carb::datasource::IDataSource* dataSource;
        carb::datasource::Connection* connection;
    };

    AssetsManager();

    ~AssetsManager();

    void notifyShutdown();

    bool getAssetAndDataSource(const char* assetUri, std::string& assetStr, DataSourceMapping& mapping);

    bool addDataSource(carb::datasource::IDataSource* dataSource,
                       carb::datasource::Connection* connection,
                       const char* prefix);
    bool removeDataSource(const char* prefix);

protected:
    DataSourceMapping m_omniDataSourceMapping = { nullptr, nullptr };
    std::unordered_map<std::string /* prefix */, DataSourceMapping> m_mapDataSources;
    std::vector<std::pair<carb::datasource::IDataSource*, carb::datasource::Connection*>> m_connections;
};

inline AssetsManager::AssetsManager()
{
    auto localConnections = getFilesystemConnections();
    for (auto& connection : localConnections)
    {
        addDataSource(std::get<1>(connection), std::get<2>(connection), std::string(1, std::get<0>(connection)).c_str());
#if CARB_PLATFORM_WINDOWS
        addDataSource(std::get<1>(connection), std::get<2>(connection),
                      std::string(1, std::tolower(std::get<0>(connection))).c_str());
#endif
        m_connections.push_back(std::make_pair(std::get<1>(connection), std::get<2>(connection)));
    }
}

inline AssetsManager::~AssetsManager()
{
    if (m_omniDataSourceMapping.dataSource && m_omniDataSourceMapping.connection)
    {
        m_omniDataSourceMapping.dataSource->disconnect(m_omniDataSourceMapping.connection);
    }

    for (auto& connection : m_connections)
    {
        connection.first->disconnect(connection.second);
    }
    m_connections.clear();
}

inline void AssetsManager::notifyShutdown()
{
    m_omniDataSourceMapping = {};
    m_connections.clear();
    m_mapDataSources.clear();
}

inline bool AssetsManager::addDataSource(carb::datasource::IDataSource* dataSource,
                                         carb::datasource::Connection* connection,
                                         const char* prefix)
{
    CARB_ASSERT(dataSource && connection);

    auto it = m_mapDataSources.find(prefix);
    if (it != m_mapDataSources.end())
    {
        if ((it->second.dataSource != dataSource) || (it->second.connection != connection))
        {
            CARB_LOG_ERROR("A different DataSource is already registered for prefix '%s'", prefix);
            return false;
        }
        return true;
    }
    DataSourceMapping& mapping = m_mapDataSources[prefix];
    mapping.dataSource = dataSource;
    mapping.connection = connection;
    CARB_LOG_INFO("Added dataSource connection for prefix '%s'", prefix);

    return true;
}

inline bool AssetsManager::removeDataSource(const char* prefix)
{
    return m_mapDataSources.erase(prefix) != 0;
}

inline bool AssetsManager::getAssetAndDataSource(const char* assetUri, std::string& assetStr, DataSourceMapping& mapping)
{
    if (assetUri == nullptr)
    {
        CARB_LOG_ERROR("assetUri is null");
        CARB_ASSERT(assetUri);
        return false;
    }

    // Extract prefix from string
    /*
        The URI should be parsed in the more accurate way according to the URI Generic Syntax.
        URI = scheme:[//authority]path[?query][#fragment]
    */
    const char* colonPos = std::strchr(assetUri, ':');
    std::string prefix;

#if CARB_PLATFORM_WINDOWS
    if (colonPos == nullptr)
    {
        CARB_LOG_ERROR("assetUri not a valid Uri: '%s'", assetUri);
        CARB_ASSERT(false);
        return false;
    }

    // Sanity check to ensure Uri contains a valid path after prefix://
    if (std::strlen(colonPos) <= 3)
    {
        CARB_LOG_ERROR("Invalid asset Uri: '%s'", assetUri);
        CARB_ASSERT(false);
        return false;
    }

    prefix = std::string(assetUri, colonPos);
#else
    // ":" on Linux implies an omniverse connection
    // clang-format off
    /*
    This comment is not exactly correct, see e.g.:

        ls /sys/bus/usb/devices/
        1-0:1.0  1-13:1.0  1-5:1.0  1-5:1.2  1-6:1.0  1-6:1.2  1-8:1.0  2-6      3-0:1.0  5-0:1.0  7-0:1.0  usb1  usb3  usb5  usb7
        1-13     1-5       1-5:1.1  1-6      1-6:1.1  1-8      2-0:1.0  2-6:1.0  4-0:1.0  6-0:1.0  8-0:1.0  usb2  usb4  usb6  usb8

    */
    // clang-format on
    if (colonPos)
    {
        prefix = std::string(assetUri, colonPos);
    }
    // on Linux the local filesystem will not contain ":" to extract a prefix from
    else
    {
        colonPos = std::strchr(assetUri, '/');

        if (colonPos == nullptr)
        {
            CARB_LOG_ERROR("assetUri not a valid Uri: '%s'", assetUri);
            CARB_ASSERT(false);
            return false;
        }

        prefix = "/";
    }
#endif

    const std::string assetUriStr = assetUri;
    bool isClientLibPath = isPathPrefixed(assetUriStr, "omniverse:", 10) || isPathPrefixed(assetUriStr, "http:", 5) || isPathPrefixed(assetUriStr, "https:", 6) || isPathPrefixed(assetUriStr, "file:", 5);
    if (isClientLibPath)
    {
        if (!m_omniDataSourceMapping.dataSource)
        {
            using namespace carb::datasource;
            IDataSource* dataSource =
                carb::getFramework()->acquireInterface<IDataSource>("carb.datasource-omniclient.plugin");

            Connection* connection = connectAndWait({ "" }, dataSource);
            if (connection)
            {
                m_omniDataSourceMapping.dataSource = dataSource;
                m_omniDataSourceMapping.connection = connection;
            }
            else
            {
                CARB_LOG_ERROR("Failed to establish omniverse connection!");
            }
        }

        mapping.dataSource = m_omniDataSourceMapping.dataSource;
        mapping.connection = m_omniDataSourceMapping.connection;
        assetStr = assetUri;
        return (mapping.dataSource != nullptr);
    }

    // Look for the prefix in our dataSource mappings
    auto it = m_mapDataSources.find(prefix);
    if (it == m_mapDataSources.end())
    {
        CARB_LOG_ERROR("assetUri '%s' contains unknown dataSource prefix '%s'", assetUri, prefix.c_str());
        return false;
    }
    else
    {
        mapping = it->second;
    }

    // Return the extracted string and dataSource mapping
    // TODO1 - HACKHACK for GTC, this is more pathing fun to handle the various ways paths can get to this code
    if (colonPos[0] == '.' && (colonPos[1] == '/' || colonPos[1] == '\\'))
    {
        // We're going to try and use this as-is
    }
    else if (colonPos[0] == '/' || colonPos[0] == '\\')
    {
        // We're just going to try and see what happens here
    }
    else
    {
        while ((colonPos[0] == ':' || colonPos[1] == ':' || colonPos[1] == '/' || colonPos[0] == '\\') &&
               colonPos[0] != '0')
        {
            colonPos++;
        }
    }
    assetStr = std::string(colonPos); // TODO???
    return true;
}

using FileSystemConnectionPair = std::pair<char, carb::datasource::Connection*>;
inline std::vector<FileSystemConnectionPair> getFileSystemConnections()
{
    carb::filesystem::IFileSystem* fs = carb::getFramework()->acquireInterface<carb::filesystem::IFileSystem>();
    carb::datasource::IDataSource* dataSourceFile = getDataSourceFile();

    std::vector<FileSystemConnectionPair> vec;

    if (fs == nullptr || dataSourceFile == nullptr)
        return vec;

#if CARB_PLATFORM_WINDOWS

    for (char drive = 'A'; drive <= 'Z'; drive++)
    {
        std::string drivePath = std::string(1, drive) + ":";
        if (fs->exists(drivePath.c_str()))
        {
            carb::datasource::Connection* connection =
                carb::datasource::connectAndWait(carb::datasource::ConnectionDesc{ drivePath.c_str() }, dataSourceFile);
            if (connection == nullptr)
            {
                CARB_LOG_ERROR("failed to get connection for drive %c", drive);
                continue;
            }

            vec.push_back(std::make_pair(drive, connection));
        }
    }

#else
    // Just add the root of the filesystem
    carb::datasource::Connection* connection =
        carb::datasource::connectAndWait(carb::datasource::ConnectionDesc{ "/" }, dataSourceFile);
    if (connection == nullptr)
        CARB_LOG_ERROR("failed to get connection for the filesystem root");
    else
        vec.push_back(std::make_pair('/', connection));
#endif

    return vec;
}
}
}
}
