#include "MctpUtil.hpp"

#include "absl/strings/match.h"

#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 <cstdint>
#include <filesystem>

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

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

struct EndpointState
{
    State state = State::None;
    bool querying = false;
    std::string emConfigPath;
    std::chrono::steady_clock::time_point lastQueryTime =
        std::chrono::steady_clock::time_point::min();
};

static std::map<std::string, EndpointState> endpointStates;
static uint8_t filterMsgType = 0;

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,
    bool bypassCooldown = false, std::function<void()> cleanup = nullptr)
{
    auto& state = endpointStates[endpointPath];
    auto now = std::chrono::steady_clock::now();
    static constexpr auto cooldown = std::chrono::seconds(5);

    state.querying = true; // Protect the window!

    // Apply rate-limiting cooldown only if this endpoint has been queried
    // before (to avoid subtracting time_point::min() which causes overflow), if
    // we are not explicitly bypassing it (e.g. for post-priming queries), and
    // if we are within the 5-second cooldown window.
    if (state.lastQueryTime != std::chrono::steady_clock::time_point::min() &&
        !bypassCooldown && (now - state.lastQueryTime < cooldown))
    {
        auto delay = cooldown - (now - state.lastQueryTime);

        auto timer =
            std::make_shared<boost::asio::steady_timer>(conn->get_io_context());
        timer->expires_after(delay);

        timer->async_wait([conn, endpointPath, emConfigPath,
                           cleanup = std::move(cleanup),
                           timer](const boost::system::error_code& ec) mutable {
            if (ec)
            {
                return;
            }
            performEmConfigQuery(conn, endpointPath, emConfigPath, false,
                                 std::move(cleanup));
        });
        return;
    }
    // Step 1: Query SupportedMessageTypes directly from mctpd
    conn->async_method_call(
        [conn, endpointPath, emConfigPath, cleanup = std::move(cleanup)](
            const boost::system::error_code& ec,
            const std::variant<std::vector<uint8_t>>& value) mutable {
        auto& state = endpointStates[endpointPath];
        if (ec)
        {
            lg2::error(
                "Failed to get SupportedMessageTypes for {PATH}: {ERROR}",
                "PATH", endpointPath, "ERROR", ec.message());
            if (cleanup)
            {
                cleanup();
            }
            else
            {
                state.querying = false;
            }
            return;
        }

        const auto* typesPtr = std::get_if<std::vector<uint8_t>>(&value);
        if (!typesPtr)
        {
            lg2::error("Invalid SupportedMessageTypes type for {PATH}", "PATH",
                       endpointPath);
            if (cleanup)
            {
                cleanup();
            }
            else
            {
                state.querying = false;
            }
            return;
        }

        // Check if it contains the requested message type
        if (std::find(typesPtr->begin(), typesPtr->end(), filterMsgType) ==
            typesPtr->end())
        {
            lg2::info(
                "Endpoint {PATH} does not support message type {TYPE}, ignoring",
                "PATH", endpointPath, "TYPE", filterMsgType);
            if (cleanup)
            {
                cleanup();
            }
            else
            {
                state.querying = false; // Clear protection!
            }
            return;
        }

        // Step 2: Proceed to EM query
        conn->async_method_call(
            [conn, endpointPath, emConfigPath, cleanup = std::move(cleanup)](
                const boost::system::error_code& ec,
                const ManagedObjectType& objects) mutable {
            auto& state = endpointStates[endpointPath];
            state.querying = false;
            state.lastQueryTime =
                std::chrono::steady_clock::now(); // Update time!

            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);
                }
                if (cleanup)
                {
                    cleanup();
                }
                return;
            }

            if (ec)
            {
                lg2::error("Failed to get managed objects for {PATH}: {ERROR}",
                           "PATH", emConfigPath, "ERROR", ec.message());
                if (cleanup)
                {
                    cleanup();
                }
                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);
                if (cleanup)
                {
                    cleanup();
                }
                return;
            }

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

            if (lastState == State::Add)
            {
                lg2::info(
                    "DEBUG: Triggering deferred query for endpoint {ENDPOINT}",
                    "ENDPOINT", endpointPath);
                performEmConfigQuery(conn, endpointPath, emConfigPath);
            }

            if (cleanup)
            {
                cleanup();
            }
        },
            "xyz.openbmc_project.EntityManager",
            "/xyz/openbmc_project/inventory",
            "org.freedesktop.DBus.ObjectManager", "GetManagedObjects");
    },
        "au.com.codeconstruct.MCTP1", endpointPath,
        "org.freedesktop.DBus.Properties", "Get",
        "xyz.openbmc_project.MCTP.Endpoint", "SupportedMessageTypes");
}

static void triggerDeferredQueries(
    const std::shared_ptr<sdbusplus::asio::connection>& conn,
    bool isPriming = false)
{
    for (auto& [path, state] : endpointStates)
    {
        if (path == "__global__")
        {
            continue;
        }

        if (state.state == State::Add)
        {
            state.state = State::None;
            performEmConfigQuery(conn, path, state.emConfigPath, isPriming);
        }
        else if (state.state == State::Remove)
        {
            state.state = State::None;
            auto it = mctpEndpointConfigMap.find(path);
            if (it != mctpEndpointConfigMap.end())
            {
                mctpEndpointConfigMap.erase(it);
            }
        }
    }
}

void setupMctpEndpointListener(
    const std::shared_ptr<sdbusplus::asio::connection>& conn,
    MctpMessageType msgType)
{
    filterMsgType = static_cast<uint8_t>(msgType);
    // Match 1: Listen for Associations from mctp-reactor
    const std::string associationMatchSpec =
        "type='signal',interface='org.freedesktop.DBus.ObjectManager',member='InterfacesAdded',arg0path='/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* endpointsPtr =
            std::get_if<std::vector<std::string>>(&propIt->second);
        if (!endpointsPtr)
        {
            return;
        }

        std::string emConfigPath = endpointsPtr->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];
        state.emConfigPath = emConfigPath;

        bool wasQuerying = endpointStates["__global__"].querying ||
                           state.querying;

        if (wasQuerying)
        {
            lg2::info("Deferring Add for endpoint {PATH}", "PATH",
                      endpointPath);
            state.state = State::Add;
        }
        else
        {
            performEmConfigQuery(conn, endpointPath, emConfigPath);
        }
    });

    const std::string associationRemoveMatchSpec =
        "type='signal',interface='org.freedesktop.DBus.ObjectManager',member='InterfacesRemoved',arg0path='/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 (endpointStates["__global__"].querying || 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
    endpointStates["__global__"].querying = true;
    auto barrier = std::make_shared<int>(0);

#ifdef UNIT_TEST
    // Inject 5 second delay ONLY for tests to hit the race condition
    auto timer =
        std::make_shared<boost::asio::steady_timer>(conn->get_io_context());
    timer->expires_after(std::chrono::seconds(5));
    timer->async_wait([conn, barrier](const boost::system::error_code& ec) {
        if (ec)
            return;
#endif

        constexpr auto associationInterfaces =
            std::to_array({"xyz.openbmc_project.Association"});
        conn->async_method_call(
            [conn, barrier](const boost::system::error_code& ec,
                            const GetSubTreeType& subtree) {
            if (ec)
            {
                lg2::error(
                    "Failed to get associations from Object Mapper: {ERROR}",
                    "ERROR", ec.message());
                endpointStates["__global__"].querying = false;
                return;
            }

            if (subtree.empty())
            {
                endpointStates["__global__"].querying = false;
                return;
            }

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

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

                (*barrier)++; // Increment barrier!

                conn->async_method_call(
                    [conn, barrier, endpointPath](
                        const boost::system::error_code& ec,
                        const std::variant<std::vector<std::string>>& value) {
                    std::weak_ptr<sdbusplus::asio::connection> weakConn = conn;
                    auto cleanup = [barrier, weakConn]() {
                        auto conn = weakConn.lock();
                        if (!conn)
                        {
                            return;
                        }
                        (*barrier)--;
                        if (*barrier == 0)
                        {
                            endpointStates["__global__"].querying = false;
                            triggerDeferredQueries(conn, true);
                        }
                    };

                    if (ec)
                    {
                        lg2::error(
                            "Failed to get endpoints for {PATH}: {ERROR}",
                            "PATH", endpointPath, "ERROR", ec.message());
                        cleanup();
                        return;
                    }

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

                    std::string emConfigPath = endpoints->front();

                    auto& state = endpointStates[endpointPath];
                    state.emConfigPath = emConfigPath;

                    performEmConfigQuery(conn, endpointPath, emConfigPath, true,
                                         std::move(cleanup));
                },
                    "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);

#ifdef UNIT_TEST
    });
#endif
}

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

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

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

                busInfo["BusType"] = busType;
            }
            else if (!busType.empty())
            {
                lg2::warning("Unsupported BusType {TYPE}", "TYPE", busType);
            }
        }
    }
    return busInfo;
}

void cleanupMctpEndpointListener()
{
    associationMatch.reset();
    associationRemoveMatch.reset();
}