#include "MctpUtil.hpp"

#include "absl/strings/match.h"

#include <boost/asio/spawn.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 <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, const std::function<void()>& cleanup,
    const boost::asio::yield_context& yield)
{
    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);

        boost::asio::steady_timer timer(conn->get_io_context());
        timer.expires_after(delay);

        boost::system::error_code ec;
        timer.async_wait(yield[ec]);
        if (ec)
        {
            if (cleanup)
            {
                cleanup();
            }
            else
            {
                state.querying = false;
            }
            return;
        }
    }
    // Step 1: Query SupportedMessageTypes directly from mctpd
    boost::system::error_code ec;
    std::variant<std::vector<uint8_t>> value;
    try
    {
        value = conn->yield_method_call<std::variant<std::vector<uint8_t>>>(
            yield, ec, "au.com.codeconstruct.MCTP1", endpointPath,
            "org.freedesktop.DBus.Properties", "Get",
            "xyz.openbmc_project.MCTP.Endpoint", "SupportedMessageTypes");
    }
    catch (const std::exception& e)
    {
        lg2::error(
            "Exception getting SupportedMessageTypes for {PATH}: {ERROR}",
            "PATH", endpointPath, "ERROR", e.what());
        if (cleanup)
        {
            cleanup();
        }
        else
        {
            state.querying = false;
        }
        return;
    }

    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 == nullptr)
    {
        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
    ManagedObjectType objects;
    try
    {
        objects = conn->yield_method_call<ManagedObjectType>(
            yield, ec, "xyz.openbmc_project.EntityManager",
            "/xyz/openbmc_project/inventory",
            "org.freedesktop.DBus.ObjectManager", "GetManagedObjects");
    }
    catch (const std::exception& e)
    {
        lg2::error("Exception getting managed objects: {ERROR}", "ERROR",
                   e.what());
        if (cleanup)
        {
            cleanup();
        }
        return;
    }

    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, false, nullptr,
                             yield);
    }

    if (cleanup)
    {
        cleanup();
    }
}

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;
            boost::asio::spawn(
                conn->get_io_context(),
                [conn, path = path, emConfigPath = state.emConfigPath,
                 isPriming](const boost::asio::yield_context& yield) {
                performEmConfigQuery(conn, path, emConfigPath, isPriming,
                                     nullptr, yield);
            });
        }
        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
        {
            boost::asio::spawn(conn->get_io_context(),
                               [conn, endpointPath, emConfigPath](
                                   const boost::asio::yield_context& yield) {
                performEmConfigQuery(conn, endpointPath, emConfigPath, false,
                                     nullptr, yield);
            });
        }
    });

    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;

                    boost::asio::spawn(
                        conn->get_io_context(),
                        [conn, endpointPath, emConfigPath,
                         cleanup = std::move(cleanup)](
                            const boost::asio::yield_context& yield) {
                        performEmConfigQuery(conn, endpointPath, emConfigPath,
                                             true, std::move(cleanup), yield);
                    });
                },
                    "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 == "I2C")
            {
                extractProp(props, busInfo, "Bus");
                extractProp(props, busInfo, "Address");

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

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