#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::map<MctpCallbackToken, MctpEndpointCallback> mctpEndpointCallbacks;
static MctpCallbackToken nextCallbackToken = 1;

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 boost::asio::yield_context& yield)
{
    auto& state = endpointStates[endpointPath];
    static constexpr auto cooldown = std::chrono::seconds(5);

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

    // Use shared_ptr constructor directly with a null pointer to act as a scope
    // guard
    std::shared_ptr<void> cleanupTrap(nullptr, [endpointPath](void*) {
        auto it = endpointStates.find(endpointPath);
        if (it != endpointStates.end())
        {
            it->second.querying = false;
        }
    });

    while (true)
    {
        auto now = std::chrono::steady_clock::now();

        // 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)
            {
                return;
            }
        }
        bypassCooldown = false;

        state.lastQueryTime =
            std::chrono::steady_clock::now(); // Update time for EVERY attempt!

        // Use a single-trip do-while loop to flatten error handling
        // and allow early exits straight to the deferred state evaluation
        do
        {
            // 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());
                break;
            }

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

            const auto* typesPtr = std::get_if<std::vector<uint8_t>>(&value);
            if (typesPtr == nullptr)
            {
                lg2::error("Invalid SupportedMessageTypes type for {PATH}",
                           "PATH", endpointPath);
                break;
            }
            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);
                break;
            }

            // 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());
                break;
            }

            if (state.state != State::Remove)
            {
                if (ec)
                {
                    lg2::error(
                        "Failed to get managed objects for {PATH}: {ERROR}",
                        "PATH", emConfigPath, "ERROR", ec.message());
                    break;
                }

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

                    for (const auto& [token, cb] : mctpEndpointCallbacks)
                    {
                        cb(endpointPath, objIt->second, false);
                    }
                }
            }
        } while (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);
                for (const auto& [token, cb] : mctpEndpointCallbacks)
                {
                    cb(endpointPath, {}, true);
                }
            }
            break;
        }
        if (lastState == State::Add)
        {
            lg2::info(
                "DEBUG: Triggering deferred query for endpoint {ENDPOINT}",
                "ENDPOINT", endpointPath);
            continue;
        }

        break;
    }
}

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

        if (state.state == State::Add)
        {
            state.state = State::None;
            performEmConfigQuery(conn, path, state.emConfigPath, isPriming,
                                 yield);
        }
        else if (state.state == State::Remove)
        {
            state.state = State::None;
            auto it = mctpEndpointConfigMap.find(path);
            if (it != mctpEndpointConfigMap.end())
            {
                mctpEndpointConfigMap.erase(it);
                for (const auto& [token, cb] : mctpEndpointCallbacks)
                {
                    cb(path, {}, true);
                }
            }
        }
    }
}

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,
                                     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);

                for (const auto& [token, cb] : mctpEndpointCallbacks)
                {
                    cb(endpointPath, {}, true);
                }
            }
        }
    });

    // Prime the cache: Query existing Associations
    endpointStates["__global__"].querying = true;

    boost::asio::spawn(conn->get_io_context(),
                       [conn](const boost::asio::yield_context& yield) {
#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));
        boost::system::error_code ec_timer;
        timer->async_wait(yield[ec_timer]);
        if (ec_timer)
            return;
#endif
        constexpr auto associationInterfaces =
            std::to_array({"xyz.openbmc_project.Association"});
        boost::system::error_code ec;
        GetSubTreeType subtree;

        try
        {
            subtree = conn->yield_method_call<GetSubTreeType>(
                yield, ec, "xyz.openbmc_project.ObjectMapper",
                "/xyz/openbmc_project/object_mapper",
                "xyz.openbmc_project.ObjectMapper", "GetSubTree",
                "/au/com/codeconstruct/mctp1", 0, associationInterfaces);
        }
        catch (const std::exception& e)
        {
            lg2::error("Exception getting managed objects: {ERROR}", "ERROR",
                       e.what());
            endpointStates["__global__"].querying = false;
            return;
        }

        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());

            std::variant<std::vector<std::string>> value;
            try
            {
                value = conn->yield_method_call<
                    std::variant<std::vector<std::string>>>(
                    yield, ec, "xyz.openbmc_project.ObjectMapper", path,
                    "org.freedesktop.DBus.Properties", "Get",
                    "xyz.openbmc_project.Association", "endpoints");
            }
            catch (const std::exception& e)
            {
                lg2::error("Failed to get endpoints for {PATH}: {ERROR}",
                           "PATH", endpointPath, "ERROR", e.what());
                continue;
            }

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

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

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

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

            performEmConfigQuery(conn, endpointPath, emConfigPath, true, yield);
        }

        endpointStates["__global__"].querying = false;
        triggerDeferredQueries(conn, true, yield);
    });
}

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();
    endpointStates.clear(); // Prevent state leaking between unit tests
}

MctpCallbackToken registerMctpEndpointCallback(MctpEndpointCallback&& cb)
{
    MctpCallbackToken token = nextCallbackToken++;
    mctpEndpointCallbacks[token] = std::move(cb);
    return token;
}

void unregisterMctpEndpointCallback(MctpCallbackToken token)
{
    mctpEndpointCallbacks.erase(token);
}

void unregisterAllMctpEndpointCallbacks()
{
    mctpEndpointCallbacks.clear();
}