#include <iostream>
#include <algorithm>
#include <stdexcept>
#include <memory>
#include <vector>
#include <sstream>
#include <map>
#include <atomic>
#include <user/senseDeviceSpec.h>
#include <user/senseApiDesc.h>
#include <xcb/xcb.h>
#include "xcbXorg.h"

/**
 * @brief Manages X server connections using XCB
 *
 * This struct manages connections to the X server using the XCB library. It
 * ensures that each unique display and screen combination has a single
 * connection, and provides RAII management for these connections.
 */
struct XcbConnection
{
    struct XConnectionIdentifier
    {
        int display;
        int screen;

        bool operator<(const XConnectionIdentifier& other) const
        {
            if (display != other.display) return display < other.display;
            return screen < other.screen;
        }

        std::string stringify() const
        {
            std::ostringstream os;
            os << "display=" << display << ", screen=" << screen;
            return os.str();
        }
    };

    /**
     * @brief RAII guard for managing X server connection lifetime
     *
     * This guard ensures proper cleanup of X server connections in case of
     * errors during attachDeviceReq. If a connection is created but the device
     * attachment fails, and the connection's refcount is 0, this guard will
     * automatically remove the connection from the connections map.
     *
     * The guard can be "committed" using commit() to indicate successful
     * device attachment, in which case it will not perform cleanup on destruction.
     */
    class ConnectionGuard
    {
        std::shared_ptr<XcbConnection> conn;
        bool committed = false;

    public:
        explicit ConnectionGuard(std::shared_ptr<XcbConnection> c)
        :   conn(std::move(c))
        {}

        void commit(void) { committed = true; }

        ~ConnectionGuard()
        {
            if (!committed && conn && conn->refCount == 0) {
                XcbConnection::connections.erase(conn->connectionIdentifier);
            }
        }
    };

    XcbConnection(const XConnectionIdentifier& id)
    :   connection(nullptr, &xcb_disconnect),
        connectionIdentifier(id), refCount(0)
    {
        // Convert to X display string format (e.g., ":0.1")
        std::string displayString = ":" + std::to_string(id.display)
            + "." + std::to_string(id.screen);

        int screenNum;
        connection.reset(xcb_connect(displayString.c_str(), &screenNum));
        if (xcb_connection_has_error(connection.get()))
        {
            throw std::runtime_error(
                std::string(__func__) + ": Failed to connect to X server "
                + connectionIdentifier.stringify());
        }

        // Verify we got the screen we asked for
        if (screenNum != id.screen)
        {
            throw std::runtime_error(
                std::string(__func__) + ": Connected to wrong screen. "
                "Requested " + connectionIdentifier.stringify()
                + " but got screen " + std::to_string(screenNum));
        }
    }

    // Delete copy/move operations - we'll manage instances through pointers
    XcbConnection(const XcbConnection&) = delete;
    XcbConnection& operator=(const XcbConnection&) = delete;
    XcbConnection(XcbConnection&&) = delete;
    XcbConnection& operator=(XcbConnection&&) = delete;

    std::unique_ptr<xcb_connection_t, decltype(&xcb_disconnect)> connection;
    XConnectionIdentifier connectionIdentifier;
    std::atomic<int> refCount;

public:
    static std::map<XConnectionIdentifier, std::shared_ptr<XcbConnection>>
        connections;

    static std::shared_ptr<XcbConnection> getOrCreateConnection(
        const XConnectionIdentifier& id)
    {
        auto it = connections.find(id);
        if (it != connections.end()) {
            return it->second;
        }

        auto conn = std::make_shared<XcbConnection>(id);
        connections.emplace(id, conn);
        return conn;
    }
};

namespace xcb_window_search {

    // Custom deleters for XCB reply types
    struct XcbReplyDeleter {
        void operator()(xcb_query_tree_reply_t* p) { free(p); }
        void operator()(xcb_get_property_reply_t* p) { free(p); }
    };

    enum class MatchType { SUBSTRING, EXACT, ID };

    static xcb_window_t findById(
        xcb_connection_t* conn, xcb_window_t root, uint32_t targetId);

    static xcb_window_t findByName(
        xcb_connection_t* conn, xcb_window_t root,
        const std::string& targetName, std::string& outWindowName,
        MatchType matchType);
}

/**
 * @brief Represents an attached device and its associated window
 *
 * This class represents a device that has been attached to the X server. It
 * manages the connection to the X server and the window selection criteria.
 */
class AttachedDevice
{
public:
    struct WindowSelector
    {
        xcb_window_search::MatchType matchType;

        XcbConnection::XConnectionIdentifier xconn;
        struct
        {
            uint32_t id;
            std::string name;
        } window;

        std::string stringify() const
        {
            std::ostringstream os;
            os << "Display: " << xconn.display
               << ", Screen: " << xconn.screen << ", Window: ";
            if (matchType == xcb_window_search::MatchType::ID) {
                os << window.id;
            } else {
                os << "\"" << window.name << "\"";
            }
            os << " (matchType="
                << (matchType == xcb_window_search::MatchType::EXACT
                    ? "exact" :
                    (matchType == xcb_window_search::MatchType::SUBSTRING)
                        ? "substring" : "id")
               << ")";
            return os.str();
        }
    };

    AttachedDevice(const hk::device::SenseDeviceSpec &spec,
                  std::shared_ptr<XcbConnection> conn)
        : deviceSpec(spec)
        // This std::move is moving ownership from a shared_ptr to a shared_ptr
        , connection(std::move(conn))
    {
        windowSelector.xconn.display = getRequiredParamAsInt(spec, "display");
        windowSelector.xconn.screen = getRequiredParamAsInt(spec, "screen");
        parseWindowSelector(spec, windowSelector);

        // Get the root window
        const xcb_setup_t* setup = xcb_get_setup(connection->connection.get());
        xcb_screen_iterator_t iter = xcb_setup_roots_iterator(setup);
        for (int i = 0; i < windowSelector.xconn.screen; ++i) {
            xcb_screen_next(&iter);
        }
        xcb_window_t root = iter.data->root;

        // Search for window
        xcb_window_t targetWindow = 0;
        if (windowSelector.matchType == xcb_window_search::MatchType::ID)
        {
            targetWindow = xcb_window_search::findById(
                connection->connection.get(), root, windowSelector.window.id);
        }
        else
        {
            targetWindow = xcb_window_search::findByName(
                connection->connection.get(), root,
                windowSelector.window.name, actualWindowName,
                windowSelector.matchType);
        }

        if (!targetWindow)
        {
            throw std::runtime_error(
                "Failed to find window "
                + (windowSelector.matchType == xcb_window_search::MatchType::ID
                    ? std::to_string(windowSelector.window.id)
                    : "\"" + windowSelector.window.name + "\"")
                + " on display " + std::to_string(windowSelector.xconn.display)
                + ", screen " + std::to_string(windowSelector.xconn.screen));
        }
    }

    hk::device::SenseDeviceSpec deviceSpec;
    WindowSelector windowSelector;
    std::shared_ptr<XcbConnection> connection;
    std::string actualWindowName;

public:
    static int getRequiredParamAsInt(
        const hk::device::SenseDeviceSpec& spec,
        const std::string& paramName)
    {
        auto it = std::find_if(
            spec.providerParams.begin(),
            spec.providerParams.end(),
            [&paramName](const auto& param) {
                return param.first == paramName;
            }
        );

        if (it == spec.providerParams.end())
        {
            throw std::runtime_error(
                "No " + paramName + " specified in provider params");
        }

        try {
            return std::stoi(it->second);
        } catch (const std::exception& e) {
            throw std::runtime_error(
                "Failed to parse '" + paramName + "' param value '"
                + it->second + "' as integer: " + e.what());
        }
    }

    static void parseWindowSelector(
        const hk::device::SenseDeviceSpec& spec,
        WindowSelector& windowSelector)
    {
        // Default match type
        windowSelector.matchType = xcb_window_search::MatchType::SUBSTRING;

        // Check if 'dev-id', 'dev-string', or 'dev-substring' is specified
        for (const auto& param : spec.apiParams)
        {
            if (param.first == "dev-id" || param.first == "devid")
            {
                windowSelector.matchType = xcb_window_search::MatchType::ID;
                break;
            }
            if (param.first == "dev-string" || param.first == "dev-str"
                || param.first == "devstr" || param.first == "devstring")
            {
                windowSelector.matchType = xcb_window_search::MatchType::EXACT;
            }
            if (param.first == "dev-substring" || param.first == "dev-substr"
                || param.first == "devsubstr" || param.first == "devsubstring")
            {
                windowSelector.matchType = xcb_window_search::
                    MatchType::SUBSTRING;
            }
        }

        if (windowSelector.matchType == xcb_window_search::MatchType::ID)
        {
            try {
                windowSelector.window.id = std::stoul(
                    spec.deviceSelector, nullptr, 0);
            } catch (const std::exception&) {
                throw std::runtime_error(
                    "Window selector: 'dev-id' present, but selector is not "
                    "numeric");
            }
        }
        else {
            windowSelector.window.name = spec.deviceSelector;
        }
    }
};

// Define the static member
std::map<XcbConnection::XConnectionIdentifier, std::shared_ptr<XcbConnection>>
    XcbConnection::connections;

static std::vector<AttachedDevice> attachedDevices;

namespace xcb_window_search {

static xcb_window_t findById(
    xcb_connection_t* conn, xcb_window_t root, uint32_t targetId)
{
    xcb_query_tree_cookie_t cookie = xcb_query_tree(conn, root);
    std::unique_ptr<xcb_query_tree_reply_t, XcbReplyDeleter> reply(
        xcb_query_tree_reply(conn, cookie, nullptr));
    if (!reply) return 0;

    // First check if current window is the target
    if (root == targetId) {
        return root;
    }

    // Then check all children
    xcb_window_t* children = xcb_query_tree_children(reply.get());
    int num_children = xcb_query_tree_children_length(reply.get());

    for (int i = 0; i < num_children; ++i)
    {
        if (children[i] == targetId) {
            return children[i];
        }

        // Recursively search this child's subtree
        if (xcb_window_t result = findById(
            conn, children[i], targetId))
        {
            return result;
        }
    }

    return 0;
}

static xcb_window_t findByName(
    xcb_connection_t* conn, xcb_window_t root,
    const std::string& targetName, std::string& outWindowName,
    MatchType matchType)
{
    xcb_query_tree_cookie_t cookie = xcb_query_tree(conn, root);
    std::unique_ptr<xcb_query_tree_reply_t, XcbReplyDeleter> reply(
        xcb_query_tree_reply(conn, cookie, nullptr));
    if (!reply) return 0;

    // First check current window
    xcb_get_property_cookie_t prop_cookie = xcb_get_property(
        conn, 0, root, XCB_ATOM_WM_NAME, XCB_ATOM_STRING, 0, 1024);

    std::unique_ptr<xcb_get_property_reply_t, XcbReplyDeleter> prop_reply(
        xcb_get_property_reply(conn, prop_cookie, nullptr));

    if (prop_reply)
    {
        int len = xcb_get_property_value_length(prop_reply.get());
        char* name = static_cast<char*>(
            xcb_get_property_value(prop_reply.get()));
        if (len > 0)
        {
            std::string windowName(name, len);
            if ((matchType == MatchType::EXACT
                && windowName == targetName)
                || (matchType == MatchType::SUBSTRING
                    && windowName.find(targetName) != std::string::npos))
            {
                outWindowName = windowName;
                return root;
            }
        }
    }

    // Then check all children
    xcb_window_t* children = xcb_query_tree_children(reply.get());
    int num_children = xcb_query_tree_children_length(reply.get());

    for (int i = 0; i < num_children; ++i)
    {
        prop_cookie = xcb_get_property(
            conn, 0, children[i],
            XCB_ATOM_WM_NAME, XCB_ATOM_STRING, 0, 1024);

        prop_reply.reset(xcb_get_property_reply(conn, prop_cookie, nullptr));
        if (prop_reply)
        {
            int len = xcb_get_property_value_length(prop_reply.get());
            char* name = static_cast<char*>(xcb_get_property_value(
                prop_reply.get()));
            if (len > 0)
            {
                std::string windowName(name, len);
                if ((matchType == MatchType::EXACT
                    && windowName == targetName)
                    || (matchType == MatchType::SUBSTRING
                        && windowName.find(targetName) != std::string::npos))
                {
                    outWindowName = windowName;
                    return children[i];
                }
            }
        }

        // Recursively search this child's subtree
        if (xcb_window_t result = findByName(
            conn, children[i], targetName, outWindowName, matchType))
        {
            return result;
        }
    }

    return 0;
}

} // namespace xcb_window_search

static hk::sense_api::sal_mlo_initializeIndFn xcbXorg_initializeInd;
static hk::sense_api::sal_mlo_finalizeIndFn xcbXorg_finalizeInd;
static hk::sense_api::sal_mlo_attachDeviceReqFn xcbXorg_attachDeviceReq;
static hk::sense_api::sal_mlo_detachDeviceReqFn xcbXorg_detachDeviceReq;

static hk::sense_api::SenseApiDesc xcbXorgApiDesc =
{
    .name = "xcb",
    .exportedImplexorApis = { { "video-implexor" } },
    .sal_mgmt_libOps = {
        .initializeInd = xcbXorg_initializeInd,
        .finalizeInd = xcbXorg_finalizeInd,
        .attachDeviceReq = xcbXorg_attachDeviceReq,
        .detachDeviceReq = xcbXorg_detachDeviceReq
    }
};

/* Below here are the exported functions that Harikoff will use both to load
 * and use this library.
 ******************************************************************************/

/**
 * @brief Get the Sense API Descriptor
 *
 * This function is used to retrieve the descriptor for the Sense API provided
 * by this library. The descriptor contains information about the API, such as
 * its name, the implexors it exports, and the management operations it
 * supports.
 *
 * @return A reference to the SenseApiDesc structure describing the API.
 */
extern HK_UNMANGLED hk::sense_api::HK_GET_SENSE_API_DESC_FN_TYPEDEF
    HK_GET_SENSE_API_DESC_FN_NAME;

const hk::sense_api::SenseApiDesc &HK_GET_SENSE_API_DESC_FN_NAME(void)
{
    return xcbXorgApiDesc;
}

int xcbXorg_initializeInd(void)
{
    return 0;
}

int xcbXorg_finalizeInd(void)
{
    XcbConnection::connections.clear();
    return 0;
}

int xcbXorg_attachDeviceReq(const hk::device::SenseDeviceSpec &desc)
{
    // Ensure connection exists before creating device. Create conn'tion if not.
    XcbConnection::XConnectionIdentifier id{
        AttachedDevice::getRequiredParamAsInt(desc, "display"),
        AttachedDevice::getRequiredParamAsInt(desc, "screen")
    };
    auto conn = XcbConnection::getOrCreateConnection(id);
    // RAII protection in case AttachDevice construction below fails
    XcbConnection::ConnectionGuard guard(conn);

    // Create device and increment connection refcount
    attachedDevices.emplace_back(desc, conn);
    // Successfully attached device, so decouple guard from RAII cleanup
    conn->refCount++;
    guard.commit();

    std::cout << "Attaching X11 window:\n  "
        << attachedDevices.back().windowSelector.stringify() << "\n"
        << "  Actual window name: \""
        << attachedDevices.back().actualWindowName << "\"\n"
        << "  Using " << (conn->refCount > 1 ? "existing" : "new")
        << " connection to X server\n";

    return 0;
}

int xcbXorg_detachDeviceReq(const hk::device::SenseDeviceSpec &spec)
{
    auto it = std::find_if(attachedDevices.begin(), attachedDevices.end(),
        [&spec](const AttachedDevice &device) {
            return device.deviceSpec == spec;
        }
    );

    if (it == attachedDevices.end())
    {
        auto displayIt = std::find_if(
            spec.providerParams.begin(), spec.providerParams.end(),
            [](const auto& param) { return param.first == "display"; });
        auto screenIt = std::find_if(
            spec.providerParams.begin(), spec.providerParams.end(),
            [](const auto& param) { return param.first == "screen"; });

        std::cerr << __func__ << ": Device not attached: "
            << "display=" << (displayIt != spec.providerParams.end()
                ? displayIt->second : "not specified")
            << ", screen=" << (screenIt != spec.providerParams.end()
                ? screenIt->second : "not specified")
            << ", selector=" << spec.deviceSelector << "\n";
        return 0;
    }

    XcbConnection::XConnectionIdentifier id{
        it->windowSelector.xconn.display,
        it->windowSelector.xconn.screen
    };

    // Atomic decrement refcount
    int newCount = --it->connection->refCount;

    // If no more references, close the connection
    if (newCount == 0)
    {
        XcbConnection::connections.erase(id);
        std::cout << "Closed X server connection (display="
            << id.display << ", screen=" << id.screen << ")\n";
    }

    attachedDevices.erase(it);
    std::cout << __func__ << ": Detached device\n";
    return 0;
}