#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"

// Key for identifying unique X server 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();
        }
    };

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

class AttachedDevice
{
public:
    struct WindowSelector
    {
        XcbConnection::XConnectionIdentifier xconn;
        struct
        {
            uint32_t id;
            std::string name;
        } window;
        bool useWindowId;

        std::string stringify() const
        {
            std::ostringstream os;
            os << "Display: " << xconn.display
               << ", Screen: " << xconn.screen
               << ", Window: " << (useWindowId ? 
                    std::to_string(window.id) :
                    "\"" + window.name + "\"");
            return os.str();
        }
    };

    AttachedDevice(const hk::device::SenseDeviceSpec &spec,
                  std::shared_ptr<XcbConnection> conn)
        : deviceSpec(spec)
        , connection(std::move(conn))
    {
        windowSelector.xconn.display = getRequiredParamAsInt(spec, "display");
        windowSelector.xconn.screen = getRequiredParamAsInt(spec, "screen");
        parseWindowSelector(spec.deviceSelector, windowSelector);
    }

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

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 std::string& selector,
        WindowSelector& windowSelector)
    {
        if (selector.length() >= 2 &&
            ((selector[0] == '"' && selector.back() == '"') ||
             (selector[0] == '\'' && selector.back() == '\'') ||
             (selector[0] == '`' && selector.back() == '`')))
        {
            windowSelector.window.name = selector.substr(
                1, selector.length() - 2);
            windowSelector.useWindowId = false;
        }
        else
        {
            try {
                windowSelector.window.id = std::stoul(selector);
                windowSelector.useWindowId = true;
            } catch (const std::exception&) {
                throw std::runtime_error(
                    "Window selector must be either a quoted string or a "
                    "numeric ID");
            }
        }
    }
};

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

static std::vector<AttachedDevice> attachedDevices;

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-xorg",
    .exportedImplexorApis = { { "video-implexor" } },
    .sal_mgmt_libOps = {
        .initializeInd = xcbXorg_initializeInd,
        .finalizeInd = xcbXorg_finalizeInd,
        .attachDeviceReq = xcbXorg_attachDeviceReq,
        .detachDeviceReq = xcbXorg_detachDeviceReq
    }
};

extern HK_UNMANGLED hk::sense_api::getSenseApiDescFn
    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)
{
    // Create temporary device to validate parameters
    AttachedDevice::WindowSelector tempSelector;

    tempSelector.xconn.display = AttachedDevice::getRequiredParamAsInt(
        desc, "display");
    tempSelector.xconn.screen = AttachedDevice::getRequiredParamAsInt(
        desc, "screen");

    // Get or create connection before constructing device
    XcbConnection::XConnectionIdentifier id{
        tempSelector.xconn.display,
        tempSelector.xconn.screen
    };
    auto conn = XcbConnection::getOrCreateConnection(id);

    // Create device with validated connection
    attachedDevices.emplace_back(desc, conn);
    conn->refCount++;

    std::cout << "Attaching X11 window:\n  "
              << attachedDevices.back().windowSelector.stringify() << "\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;
}