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

namespace xcb_xorg {

// Static member initialization
std::map<ConnectionIdentifier, std::shared_ptr<XcbConnection>>
    ConnectionManager::connections;

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

XcbConnection::XcbConnection(const ConnectionIdentifier& 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));
    }
}

std::shared_ptr<XcbConnection> ConnectionManager::getOrCreateConnection(
    const ConnectionIdentifier& 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;
}

void ConnectionManager::cleanupAllConnections()
{
    connections.clear();
}

size_t ConnectionManager::getConnectionCount()
{
    return connections.size();
}

/**
 * @brief Remove a specific connection from the manager
 * @param id Connection identifier to remove
 */
void ConnectionManager::removeConnection(const ConnectionIdentifier& id)
{
    auto it = connections.find(id);
    if (it != connections.end()) {
        connections.erase(it);
    }
}

namespace window_search {

struct XcbReplyDeleter {
    void operator()(xcb_query_tree_reply_t* p) { free(p); }
    void operator()(xcb_get_property_reply_t* p) { free(p); }
};

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;

    if (root == targetId) return root;

    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;
}

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 window_search

} // namespace xcb_xorg

// Export functions for external libraries
/**
 * @brief Get or create a connection to the X server
 * @param display Display number
 * @param screen Screen number
 * @return Shared pointer to connection (as void* for C compatibility)
 */
extern "C" get_or_create_connection_fn xcb_xorg_get_or_create_connection;
std::shared_ptr<xcb_xorg::XcbConnection> xcb_xorg_get_or_create_connection(
    int display, int screen)
{
    xcb_xorg::ConnectionIdentifier id{display, screen};
    auto conn = xcb_xorg::ConnectionManager::getOrCreateConnection(id);
    conn->incrementRefCount();
    return conn;
}

/**
 * @brief Clean up all connections
 */
extern "C" cleanup_connections_fn xcb_xorg_cleanup_connections;
void xcb_xorg_cleanup_connections()
{
    xcb_xorg::ConnectionManager::cleanupAllConnections();
}

/**
 * @brief Get the number of active connections
 * @return Number of active connections
 */
size_t xcb_xorg_get_connection_count()
{
    return xcb_xorg::ConnectionManager::getConnectionCount();
}

/**
 * @brief Find window by ID
 * @param conn Connection pointer (from xcb_xorg_get_connection)
 * @param root Root window
 * @param targetId Target window ID
 * @return Window ID if found, 0 if not found
 */
extern "C" find_window_by_id_fn xcb_xorg_find_window_by_id;
xcb_window_t xcb_xorg_find_window_by_id(void* conn, xcb_window_t root, uint32_t targetId)
{
    if (!conn) return 0;
    auto connection = static_cast<xcb_xorg::XcbConnection*>(conn);
    return xcb_xorg::window_search::findById(
        connection->getConnection(), root, targetId);
}

/**
 * @brief Find window by name
 * @param conn Connection pointer (from xcb_xorg_get_connection)
 * @param root Root window
 * @param targetName Target window name
 * @param outWindowName Output parameter for actual window name
 * @param matchType Type of name matching
 * @return Window ID if found, 0 if not found
 */
extern "C" find_window_by_name_fn xcb_xorg_find_window_by_name;
xcb_window_t xcb_xorg_find_window_by_name(void* conn, xcb_window_t root,
    const std::string& targetName, std::string& outWindowName,
    xcb_xorg::window_search::MatchType matchType)
{
    if (!conn) return 0;
    auto connection = static_cast<xcb_xorg::XcbConnection*>(conn);
    return xcb_xorg::window_search::findByName(
        connection->getConnection(), root, targetName, outWindowName,
        matchType);
}

/**
 * @brief Dereference a connection (decrements ref count and closes if zero)
 * @param conn Shared pointer to the connection to dereference
 */
extern "C" dereference_connection_fn xcb_xorg_dereference_connection;
void xcb_xorg_dereference_connection(std::shared_ptr<xcb_xorg::XcbConnection> conn)
{
    if (!conn) return;

    int newRefCount = conn->decrementRefCount();
    // Remove from connection manager if ref count reaches zero
    if (newRefCount <= 0) {
        xcb_xorg::ConnectionManager::removeConnection(conn->getIdentifier());
    }
}