/*
// Copyright (c) 2019 Intel Corporation
//
// 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 "absl/strings/match.h"

#include "MctpEndpoint.hpp"
#include "MctpReactorDevice.hpp"
#include "MctpUtil.hpp"
#include "NVMeBasic.hpp"
#include "NVMeDevice.hpp"
#include "NVMeIntf.hpp"
#include "NVMeMi.hpp"
#include "NVMePlugin.hpp"
#include "NVMeSubsys.hpp"
#include "Utils.hpp"

#include <dlfcn.h>

#include <boost/algorithm/string.hpp>
#include <boost/asio/steady_timer.hpp>
#include <phosphor-logging/lg2.hpp>

#include <filesystem>
#include <fstream>
#include <optional>
#include <regex>
#include <system_error>
#include <unordered_set>

// a map with key value of {path, NVMeSubsystem}
using NVMeSubsystemMap = std::map<std::string, std::shared_ptr<NVMeSubsystem>>;
static NVMeSubsystemMap nvmeSubsystems;

// a map with key value of {path, NVMeDevice}
using NVMeDeviceMap = std::map<std::string, std::shared_ptr<NVMeDevice>>;
static NVMeDeviceMap nvmeDevices;

// A map from root bus number to the Worker
// This map means to reuse the same worker for all NVMe EP under the same
// I2C root bus. There is no real physical concurrency among the i2c/mctp
// devices on the same bus. Though mctp kernel drive can schedule and
// sequencialize the transactions but assigning individual worker thread to
// each EP makes no sense.
static std::map<int, std::weak_ptr<NVMeMiWorker>> i2cWorkerMap{};

std::unordered_map<std::string, void*> pluginLibMap = {};

static std::unordered_set<int> bannedBuses;

static void initBannedI2cBus()
{
    const std::string script = "/usr/bin/init-banned-i2c-bus.sh";
    const std::string confPath = "/var/run/nvmed/banned-i2c-bus.conf";
    if (!std::filesystem::exists(script))
    {
        lg2::error("Script {SCRIPT} doesn't exist", "SCRIPT", script);
        return;
    }

    lg2::info("Begin to execute {SCRIPT}", "SCRIPT", script);
    // TODO: (b/376270522) cert-env33-c error: calling 'system' uses a command
    // processor
    int rc = std::system(script.c_str()); // NOLINT(cert-env33-c)
    lg2::info("Shell script rc = {RC}", "RC", rc);

    if (!std::filesystem::exists(confPath))
    {
        lg2::warning("Warning: {CONF_PATH} doesn't exist.", "CONF_PATH",
                     confPath);
        // Be optimistic，assume no bus is banned
        return;
    }
    std::ifstream file;
    file.open(confPath);
    if (!file.is_open())
    {
        lg2::error("Error: cannot open {CONF_PATH}", "CONF_PATH", confPath);
        // Be optimistic，assume no bus is banned
        return;
    }

    bannedBuses.clear();
    int i2cBus{0};
    while (file >> i2cBus)
    {
        lg2::info("Banned i2c bus: {BUS}", "BUS", i2cBus);
        bannedBuses.insert(i2cBus);
    }
    file.close();
}

static std::optional<int>
    extractBusNumber(const std::string& path,
                     const SensorBaseConfigMap& properties)
{
    auto findBus = properties.find("Bus");
    if (findBus == properties.end())
    {
        lg2::error("could not determine bus number for '{PATH}'", "PATH", path);
        return std::nullopt;
    }

    return std::visit(VariantToIntVisitor(), findBus->second);
}

static std::optional<int> extractAddress(const std::string& path,
                                         const SensorBaseConfigMap& properties)
{
    auto findAddr = properties.find("Address");
    if (findAddr == properties.end())
    {
        lg2::error("could not determine address for '{PATH}'", "PATH", path);
        return std::nullopt;
    }

    return std::visit(VariantToIntVisitor(), findAddr->second);
}

static std::optional<std::string>
    extractName(const std::string& path, const SensorBaseConfigMap& properties)
{
    auto findName = properties.find("Name");
    if (findName == properties.end())
    {
        lg2::error("could not determine configuration name for '{PATH}'",
                   "PATH", path);
        return std::nullopt;
    }

    return std::get<std::string>(findName->second);
}

static std::optional<std::string>
    extractProtocol(const std::string& path,
                    const SensorBaseConfigMap& properties)
{
    auto findProtocol = properties.find("Protocol");
    if (findProtocol == properties.end())
    {
        lg2::error("could not determine nvme protocol for '{PATH}'", "PATH",
                   path);
        return std::nullopt;
    }
    return std::get<std::string>(findProtocol->second);
}

static std::optional<std::string>
    extractMctpReactorConfigPath(const SensorBaseConfigMap& properties)
{
    auto findMctpReactorConfigPath = properties.find("MctpReactorConfigPath");
    if (findMctpReactorConfigPath == properties.end())
    {
        return std::nullopt;
    }
    return std::get<std::string>(findMctpReactorConfigPath->second);
}

static bool hasLockdownFeature(const std::string& path,
                               const SensorBaseConfigMap& properties)
{
    auto it = properties.find("SupportedFeatures");
    if (it == properties.end())
    {
        return false;
    }

    const auto& value = it->second;

    if (std::holds_alternative<std::vector<std::string>>(value))
    {
        const auto& features =
            std::get<std::vector<std::string>>(value); // Safe now
        for (const auto& feature : features)
        {
            if (feature == "Lockdown")
            {
                return true;
            }
        }
    }
    else
    {
        lg2::warning(
            "'{PATH}': 'SupportedFeatures' is not an array of strings.", "PATH",
            path);
    }
    return false;
}

static std::vector<BusInfo> extractNVMeBusInfo(const SensorData& configData)
{
    std::vector<BusInfo> busInfoList;
    std::map<int, BusInfo> indexedBusInfo;
    BusInfo legacyBusInfo;
    bool hasLegacy = false;

    for (const auto& pair : configData)
    {
        const std::string& intf = pair.first;
        const SensorBaseConfigMap& props = pair.second;
        if (absl::StrContains(
                intf, "xyz.openbmc_project.Configuration.NVME1000.BusInfo"))
        {
            BusInfo busInfo;
            std::string busType;
            auto it = props.find("BusType");
            if (it != props.end())
            {
                busType = std::visit(VariantToStringVisitor(), it->second);
                busInfo["BusType"] = busType;
            }

            auto extractProp = [&](const std::string& key) {
                auto it = props.find(key);
                if (it != props.end())
                {
                    busInfo[key] = std::visit(VariantToStringVisitor(),
                                              it->second);
                }
            };

            if (busType == "USB")
            {
                extractProp("RootHubPath");
                extractProp("Port");
                extractProp("InterfaceNum");
                extractProp("Configuration");

                if (busInfo.find("Configuration") == busInfo.end())
                {
                    busInfo["Configuration"] = "1";
                }

                if (busInfo.find("InterfaceNum") == busInfo.end())
                {
                    busInfo["InterfaceNum"] = "0";
                }
            }
            else if (busType == "I2C")
            {
                extractProp("Bus");
                extractProp("Address");
            }
            else if (busType == "MCTP Bridge")
            {
                extractProp("EIDOffset");
            }
            else if (!busType.empty())
            {
                lg2::warning("Unsupported BusType {TYPE} in NVME1000 config",
                             "TYPE", busType);
            }

            std::string suffix = intf.substr(
                std::string(
                    "xyz.openbmc_project.Configuration.NVME1000.BusInfo")
                    .length());
            if (suffix.empty())
            {
                legacyBusInfo = busInfo;
                hasLegacy = true;
            }
            else
            {
                try
                {
                    int index = std::stoi(suffix);
                    indexedBusInfo[index] = busInfo;
                }
                catch (const std::exception& e)
                {
                    lg2::error(
                        "Failed to parse index from interface {INTF}: {ERROR}",
                        "INTF", intf, "ERROR", e.what());
                }
            }
        }
    }

    if (hasLegacy)
    {
        busInfoList.push_back(legacyBusInfo);
    }
    else
    {
        for (const auto& [index, busInfo] : indexedBusInfo)
        {
            busInfoList.push_back(busInfo);
        }
    }

    return busInfoList;
}

static void handleConfigurations(
    boost::asio::io_context& io, sdbusplus::asio::object_server& objectServer,
    std::shared_ptr<sdbusplus::asio::connection>& dbusConnection,
    const ManagedObjectType& nvmeConfigurations)
{
    // Initialize banned i2c bus info on every configuration change
    initBannedI2cBus();

    for (const auto& [nvmeObjectPath, configData] : nvmeConfigurations)
    {
        // find base configuration
        auto sensorBase =
            configData.find(configInterfaceName(nvme::sensorType));
        if (sensorBase == configData.end())
        {
            continue;
        }

        const SensorBaseConfigMap& sensorConfig = sensorBase->second;
        std::optional<std::string> nvmeProtocol =
            extractProtocol(nvmeObjectPath, sensorConfig);
        if (!nvmeProtocol)
        {
            nvmeProtocol.emplace("mi_basic");
        }

        std::optional<int> busNumber;
        std::optional<int> address;
        std::vector<BusInfo> busInfoList;
        std::string endpointPath;

        if (*nvmeProtocol == "mi_basic" || *nvmeProtocol == "mi_i2c")
        {
            busNumber = extractBusNumber(nvmeObjectPath, sensorConfig);
            address = extractAddress(nvmeObjectPath, sensorConfig);
            if (!busNumber)
            {
                continue;
            }
        }
        else if (*nvmeProtocol == "mi_mctp")
        {
            busInfoList = extractNVMeBusInfo(configData);
            if (busInfoList.empty())
            {
                lg2::error("Missing BusInfo for mi_mctp device {PATH}", "PATH",
                           nvmeObjectPath.str);
                continue;
            }

            // Matching is now handled in MctpReactorDevice::setup()

            if (!busInfoList.empty())
            {
                const auto& rootBusInfo = busInfoList.front();
                auto itType = rootBusInfo.find("BusType");
                if (itType != rootBusInfo.end() && itType->second == "I2C")
                {
                    auto itBus = rootBusInfo.find("Bus");
                    if (itBus != rootBusInfo.end())
                    {
                        busNumber = std::stoi(itBus->second);
                    }
                    auto itAddr = rootBusInfo.find("Address");
                    if (itAddr != rootBusInfo.end())
                    {
                        address = std::stoi(itAddr->second, nullptr, 0);
                    }
                }
            }
        }

        std::optional<std::string> sensorName = extractName(nvmeObjectPath,
                                                            sensorConfig);
        std::optional<std::string> mctpReactorConfigPath =
            extractMctpReactorConfigPath(sensorConfig);

        bool enableFeatureLockdown = hasLockdownFeature(nvmeObjectPath,
                                                        sensorConfig);

        if (!sensorName)
        {
            continue;
        }

        const bool isMctpReactorDevice = mctpReactorConfigPath.has_value() ||
                                         *nvmeProtocol == "mi_mctp";

        if (busNumber && bannedBuses.contains(*busNumber))
        {
            lg2::info("Skip banned i2c bus: {BUS}", "BUS", *busNumber);
            continue;
        }

        if (*nvmeProtocol == "mi_basic")
        {
            // defualt i2c basic port is 0x6a
            if (!address)
            {
                address.emplace(0x6a);
            }
            try
            {
                NVMeIntf nvmeIntf = NVMeIntf::create<NVMeBasic>(io, *busNumber,
                                                                *address);

                auto nvmeSubsys = NVMeSubsystem::create(
                    io, objectServer, dbusConnection, nvmeObjectPath,
                    *sensorName, configData, nvmeIntf, enableFeatureLockdown);
                auto nvmeDev = NVMeDevice::create(*sensorName, nvmeIntf);

                auto [subsysEntry, subsysAdded] =
                    nvmeSubsystems.try_emplace(nvmeObjectPath, nvmeSubsys);
                if (subsysAdded)
                {
                    subsysEntry->second->start();
                }

                auto [devEntry, devAdded] =
                    nvmeDevices.try_emplace(nvmeObjectPath, nvmeDev);
                if (devAdded)
                {
                    devEntry->second->start();
                }
            }
            catch (std::exception& ex)
            {
                lg2::error(
                    "Failed to add nvme basic interface for '{PATH}': {ERROR}",
                    "PATH", nvmeObjectPath.str, "ERROR", ex.what());
            }
        }
        else if (*nvmeProtocol == "mi_i2c" || *nvmeProtocol == "mi_mctp")
        {
            // defualt i2c nvme-mi port is 0x1d
            if (!address)
            {
                address.emplace(0x1d);
            }
            PowerState powerState = getPowerState(sensorConfig);

            std::shared_ptr<NVMeMiWorker> worker;
            if (singleWorkerFeature && busNumber.has_value())
            {
                auto root = deriveRootBus(*busNumber);

                if (!root || *root < 0)
                {
                    throw std::runtime_error("invalid root bus number");
                }
                auto res = i2cWorkerMap.find(*root);

                if (res == i2cWorkerMap.end() || res->second.expired())
                {
                    worker = NVMeMiWorker::create(io);
                    i2cWorkerMap[*root] = worker;
                }
                else
                {
                    worker = res->second.lock();
                }
            }
            else
            {
                worker = NVMeMiWorker::create(io);
            }

            try
            {
                std::shared_ptr<MctpDevice> mctpDev;
                if (*nvmeProtocol == "mi_mctp")
                {
                    mctpDev = std::make_shared<MctpReactorDevice>(
                        dbusConnection, nvmeObjectPath.str, busInfoList,
                        busNumber, address);
                }
                else if (!isMctpReactorDevice)
                {
                    mctpDev = std::make_shared<SmbusMctpdDevice>(
                        dbusConnection, *busNumber, *address);
                }
                else // isMctpReactorDevice
                {
                    mctpDev = std::make_shared<MctpReactorDevice>(
                        dbusConnection, *mctpReactorConfigPath, busNumber,
                        address);
                }

                NVMeIntf nvmeIntf = NVMeIntf::create<NVMeMi>(
                    io, dbusConnection, mctpDev, worker, powerState);
                auto nvmeSubsys = NVMeSubsystem::create(
                    io, objectServer, dbusConnection, nvmeObjectPath,
                    *sensorName, configData, nvmeIntf, enableFeatureLockdown);

                auto nvmeDev = NVMeDevice::create(io, *sensorName, mctpDev,
                                                  nvmeIntf);

                auto [subsysEntry, subsysAdded] =
                    nvmeSubsystems.try_emplace(nvmeObjectPath, nvmeSubsys);
                if (subsysAdded)
                {
                    subsysEntry->second->start();
                }

                auto [devEntry, devAdded] =
                    nvmeDevices.try_emplace(nvmeObjectPath, nvmeDev);
                if (devAdded)
                {
                    devEntry->second->start();
                }
            }
            catch (std::exception& ex)
            {
                lg2::error(
                    "Failed to add nvme mi interface for '{PATH}': {ERROR}",
                    "PATH", nvmeObjectPath.str, "ERROR", ex.what());
            }
        }
    }
}

void createNVMeSubsystems(
    boost::asio::io_context& io, sdbusplus::asio::object_server& objectServer,
    std::shared_ptr<sdbusplus::asio::connection>& dbusConnection)
{
    // todo: it'd be better to only update the ones we care about
    for (auto& [_, nvmeDev] : nvmeDevices)
    {
        nvmeDev->stop();
    }
    nvmeDevices.clear();

    for (auto& [_, nvmeSubsys] : nvmeSubsystems)
    {
        nvmeSubsys->stop();
    }
    nvmeSubsystems.clear();

    static int count = 0;
    static ManagedObjectType configs;
    count += 2;

    auto getter = std::make_shared<GetSensorConfiguration>(
        dbusConnection, [&io, &objectServer, &dbusConnection](
                            const ManagedObjectType& nvmeConfigurations) {
        configs = nvmeConfigurations;
        count--;
        if (count == 0)
        {
            handleConfigurations(io, objectServer, dbusConnection, configs);
        }
        else
        {
            lg2::error("more than one `handleConfigurations` has been "
                       "scheduled, cancel the current one");
        }
    });
    auto timer = std::make_shared<boost::asio::steady_timer>(
        io, std::chrono::seconds(10));
    timer->async_wait([&io, &objectServer, &dbusConnection,
                       timer](const boost::system::error_code& ec) {
        count--;
        if (ec)
        {
            return;
        }
        if (count == 0)
        {
            handleConfigurations(io, objectServer, dbusConnection, configs);
        }
        else
        {
            lg2::error("`handleConfigurations` has not been triggered, "
                       "cancel the timer");
        }
    });

    getter->getConfiguration(std::vector<std::string>{nvme::sensorType});
}

static void interfaceRemoved(sdbusplus::message_t& message)
{
    if (message.is_method_error())
    {
        lg2::error("interfacesRemoved callback method error");
        return;
    }

    sdbusplus::message::object_path path;
    std::vector<std::string> interfaces;

    message.read(path, interfaces);

    auto interface = std::find(interfaces.begin(), interfaces.end(),
                               configInterfaceName(nvme::sensorType));
    if (interface == interfaces.end())
    {
        return;
    }

    auto device = nvmeDevices.find(path);
    if (device != nvmeDevices.end())
    {
        device->second->stop();
        nvmeDevices.erase(device);
    }

    auto subsys = nvmeSubsystems.find(path);
    if (subsys != nvmeSubsystems.end())
    {
        subsys->second->stop();
        nvmeSubsystems.erase(subsys);
    }
}

int main(int argc, char** argv)
{
    for (int i = 1; i < argc; ++i)
    {
        std::string_view arg(argv[i]);
        if (arg == "--enable-metrics")
        {
            setMetricStoreEnabled(true);
            lg2::info("Metric store runtime override: enabled via CLI");
        }
        else if (arg == "--disable-metrics")
        {
            setMetricStoreEnabled(false);
            lg2::info("Metric store runtime override: disabled via CLI");
        }
    }

    if (singleWorkerFeature)
    {
        lg2::info("singleWorkerFeature on");
    }

    // Load plugin shared libraries
    try
    {
        for (const auto& entry :
             std::filesystem::directory_iterator(NVMePlugin::libraryPath))
        {
            void* lib = dlopen(entry.path().c_str(), RTLD_NOW);
            if (lib != nullptr)
            {
                pluginLibMap.emplace(entry.path().filename().string(), lib);
            }
            else
            {
                lg2::error("could not load the plugin: {ERROR}", "ERROR",
                           dlerror());
            }
        }
    }
    catch (const std::filesystem::filesystem_error& e)
    {
        lg2::error("failed to open plugin folder: {ERROR}", "ERROR", e.what());
    }

    // TODO: set single thread mode according to input parameters

    boost::asio::io_context io;
    auto systemBus = std::make_shared<sdbusplus::asio::connection>(io);
    systemBus->request_name("xyz.openbmc_project.NVMe");
    sdbusplus::asio::object_server objectServer(systemBus, true);
    objectServer.add_manager("/xyz/openbmc_project/sensors");
    objectServer.add_manager("/xyz/openbmc_project/inventory");

    setupMctpEndpointListener(systemBus, MctpMessageType::NVME_MI);

    boost::asio::post(
        io, [&]() { createNVMeSubsystems(io, objectServer, systemBus); });

    boost::asio::steady_timer filterTimer(io);
    std::function<void(sdbusplus::message_t&)> eventHandler =
        [&filterTimer, &io, &objectServer, &systemBus](sdbusplus::message_t&) {
        // this implicitly cancels the timer
        filterTimer.expires_after(std::chrono::seconds(1));

        filterTimer.async_wait([&](const boost::system::error_code& ec) {
            if (ec == boost::asio::error::operation_aborted)
            {
                return; // we're being canceled
            }

            if (ec)
            {
                lg2::error("Error: {ERROR}", "ERROR", ec.message());
                return;
            }

            createNVMeSubsystems(io, objectServer, systemBus);
        });
    };

    std::vector<std::unique_ptr<sdbusplus::bus::match_t>> matches =
        setupPropertiesChangedMatches(
            *systemBus, std::to_array<const char*>({NVMeSensor::sensorType}),
            eventHandler);

    // Watch for entity-manager to remove configuration interfaces
    // so the corresponding sensors can be removed.
    auto ifaceRemovedMatch = std::make_unique<sdbusplus::bus::match_t>(
        static_cast<sdbusplus::bus_t&>(*systemBus),
        "type='signal',member='InterfacesRemoved',arg0path='" +
            std::string(inventoryPath) + "/'",
        [](sdbusplus::message_t& msg) { interfaceRemoved(msg); });

    setupManufacturingModeMatch(*systemBus);

    // The NVMe controller used pipe to transfer raw data. The pipe could be
    // closed by the client. It should not be considered as an error.
    boost::asio::signal_set signals(io, SIGPIPE);
    signals.async_wait(
        [](const boost::system::error_code& error, int signalNumber) {
        lg2::info("signal: {SIGNAL}, {ERROR}", "SIGNAL",
                  strsignal(signalNumber), "ERROR", error.message());
    });
    io.run();

    for (const auto& [_, lib] : pluginLibMap)
    {
        dlclose(lib);
    }
}