#include "MctpUtil.hpp"

#include "Utils.hpp"
#include "VariantVisitors.hpp"

#include <boost/asio/steady_timer.hpp>
#include <boost/container/flat_map.hpp>
#include <phosphor-logging/lg2.hpp>
#include <sdbusplus/bus/match.hpp>

#include <chrono>
#include <filesystem>
#include <iostream>

std::map<std::string, SensorData> mctpEndpointConfigMap;

enum class State : uint8_t
{
    None,
    Add,
    Remove
};

struct EndpointState
{
    State state = State::None;
    bool querying = false;
};

static std::map<std::string, EndpointState> endpointStates;

static std::unique_ptr<sdbusplus::bus::match_t> associationMatch = nullptr;
static std::unique_ptr<sdbusplus::bus::match_t> associationRemoveMatch =
    nullptr;

static void performEmConfigQuery(
    const std::shared_ptr<sdbusplus::asio::connection>& conn,
    const std::string& endpointPath, const std::string& emConfigPath)
{
    std::filesystem::path p(emConfigPath);
    std::string parentPath = p.parent_path().string();

    conn->async_method_call(
        [conn, endpointPath, emConfigPath](const boost::system::error_code& ec,
                                           const ManagedObjectType& objects) {
        auto& state = endpointStates[endpointPath];
        state.querying = false;

        State lastState = state.state;
        state.state = State::None; // Reset for next query

        if (lastState == State::Remove)
        {
            lg2::info("Ignoring query reply for removed endpoint {ENDPOINT}",
                      "ENDPOINT", endpointPath);
            auto it = mctpEndpointConfigMap.find(endpointPath);
            if (it != mctpEndpointConfigMap.end())
            {
                mctpEndpointConfigMap.erase(it);
            }
            return;
        }

        if (ec)
        {
            lg2::error("Failed to get managed objects for {PATH}: {ERROR}",
                       "PATH", emConfigPath, "ERROR", ec.message());
            return;
        }

        auto objIt =
            objects.find(sdbusplus::message::object_path(emConfigPath));
        if (objIt == objects.end())
        {
            lg2::warning("EM config path {PATH} not found in managed objects",
                         "PATH", emConfigPath);
            return;
        }

        if (!objIt->second.empty())
        {
            lg2::info(
                "Recorded MCTP endpoint {ENDPOINT} with config from {PATH}",
                "ENDPOINT", endpointPath, "PATH", emConfigPath);
            mctpEndpointConfigMap[endpointPath] = objIt->second;
        }

        if (lastState == State::Add)
        {
            state.querying = true;
            performEmConfigQuery(conn, endpointPath, emConfigPath);
        }
    },
        "xyz.openbmc_project.EntityManager", parentPath,
        "org.freedesktop.DBus.ObjectManager", "GetManagedObjects");
}

void setupMctpEndpointListener(
    const std::shared_ptr<sdbusplus::asio::connection>& conn)
{
    // Match 1: Listen for Associations from mctp-reactor
    const std::string associationMatchSpec =
        "type='signal',interface='org.freedesktop.DBus.ObjectManager',member='InterfacesAdded',path_namespace='/au/com/codeconstruct/mctp1'";

    associationMatch = std::make_unique<sdbusplus::bus::match_t>(
        static_cast<sdbusplus::bus_t&>(*conn), associationMatchSpec,
        [conn](sdbusplus::message_t& msg) {
        sdbusplus::message::object_path path;
        boost::container::flat_map<
            std::string,
            boost::container::flat_map<
                std::string,
                std::variant<std::string, std::vector<std::string>>>>
            interfaces;
        msg.read(path, interfaces);

        auto it = interfaces.find("xyz.openbmc_project.Association");
        if (it == interfaces.end())
        {
            return;
        }

        auto propIt = it->second.find("endpoints");
        if (propIt == it->second.end())
        {
            return;
        }

        const auto* endpoints =
            std::get_if<std::vector<std::string>>(&propIt->second);
        if (!endpoints || endpoints->empty())
        {
            return;
        }

        std::string emConfigPath = endpoints->front();
        std::string endpointPath = path.str;
        constexpr std::string_view suffix = "/configured_by";
        if (endpointPath.ends_with("/configured_by"))
        {
            endpointPath =
                endpointPath.substr(0, endpointPath.size() - suffix.length());
        }

        auto& state = endpointStates[endpointPath];

        if (state.querying)
        {
            state.state = State::Add;
            return;
        }

        state.querying = true;
        performEmConfigQuery(conn, endpointPath, emConfigPath);
    });

    // Match 2: Listen for Associations (Removed)
    const std::string associationRemoveMatchSpec =
        "type='signal',interface='org.freedesktop.DBus.ObjectManager',member='InterfacesRemoved',path_namespace='/au/com/codeconstruct/mctp1'";

    associationRemoveMatch = std::make_unique<sdbusplus::bus::match_t>(
        static_cast<sdbusplus::bus_t&>(*conn), associationRemoveMatchSpec,
        [](sdbusplus::message_t& msg) {
        sdbusplus::message::object_path path;
        std::vector<std::string> interfaces;
        msg.read(path, interfaces);

        std::string endpointPath = path.str;
        constexpr std::string_view suffix = "/configured_by";
        if (endpointPath.ends_with("/configured_by"))
        {
            endpointPath =
                endpointPath.substr(0, endpointPath.size() - suffix.length());
        }

        auto& state = endpointStates[endpointPath];

        if (state.querying)
        {
            state.state = State::Remove;
        }
        else
        {
            auto it = mctpEndpointConfigMap.find(endpointPath);
            if (it != mctpEndpointConfigMap.end())
            {
                lg2::info("Removing MCTP endpoint {ENDPOINT} from map",
                          "ENDPOINT", endpointPath);
                mctpEndpointConfigMap.erase(it);
            }
        }
    });

    // Prime the cache: Query existing Associations
    constexpr auto associationInterfaces =
        std::to_array({"xyz.openbmc_project.Association"});
    conn->async_method_call(
        [conn](const boost::system::error_code& ec,
               const GetSubTreeType& subtree) {
        if (ec)
        {
            lg2::error("Failed to get associations from Object Mapper: {ERROR}",
                       "ERROR", ec.message());
            return;
        }

        for (const auto& [path, services] : subtree)
        {
            if (!path.ends_with("/configured_by"))
            {
                continue;
            }

            conn->async_method_call(
                [conn,
                 path](const boost::system::error_code& ec,
                       const std::variant<std::vector<std::string>>& value) {
                if (ec)
                {
                    lg2::error("Failed to get endpoints for {PATH}: {ERROR}",
                               "PATH", path, "ERROR", ec.message());
                    return;
                }

                const auto* endpoints =
                    std::get_if<std::vector<std::string>>(&value);
                if (!endpoints || endpoints->empty())
                {
                    return;
                }

                std::string emConfigPath = endpoints->front();
                std::string endpointPath = path;
                constexpr std::string_view suffix = "/configured_by";
                if (endpointPath.ends_with("/configured_by"))
                {
                    endpointPath = endpointPath.substr(0, endpointPath.size() -
                                                              suffix.length());
                }

                // Now query the EM config path for its properties
                std::filesystem::path p(emConfigPath);
                std::string parentPath = p.parent_path().string();

                conn->async_method_call(
                    [endpointPath,
                     emConfigPath](const boost::system::error_code& ec,
                                   const ManagedObjectType& objects) {
                    if (ec)
                    {
                        lg2::error(
                            "Failed to get managed objects for {PATH}: {ERROR}",
                            "PATH", emConfigPath, "ERROR", ec.message());
                        return;
                    }

                    auto objIt = objects.find(
                        sdbusplus::message::object_path(emConfigPath));
                    if (objIt == objects.end())
                    {
                        lg2::warning(
                            "EM config path {PATH} not found in managed objects",
                            "PATH", emConfigPath);
                        return;
                    }

                    if (!objIt->second.empty())
                    {
                        lg2::info(
                            "Primed match for endpoint {ENDPOINT} with config from {PATH}",
                            "ENDPOINT", endpointPath, "PATH", emConfigPath);
                        mctpEndpointConfigMap[endpointPath] = objIt->second;
                    }
                },
                    "xyz.openbmc_project.EntityManager", parentPath,
                    "org.freedesktop.DBus.ObjectManager", "GetManagedObjects");
            },
                "xyz.openbmc_project.ObjectMapper", path,
                "org.freedesktop.DBus.Properties", "Get",
                "xyz.openbmc_project.Association", "endpoints");
        }
    },
        "xyz.openbmc_project.ObjectMapper",
        "/xyz/openbmc_project/object_mapper",
        "xyz.openbmc_project.ObjectMapper", "GetSubTree",
        "/au/com/codeconstruct/mctp1", 0, associationInterfaces);
}