// Copyright (C) 2026 UnionTech Software Technology Co., Ltd.
// SPDX-License-Identifier: Apache-2.0 OR LGPL-3.0-only OR GPL-2.0-only OR GPL-3.0-only

#include "keyboardstatenotifymanagerinterfacev1.h"
#include "qwayland-server-treeland-keyboard-state-notify-unstable-v1.h"
#include "seat/helper.h"
#include "seatsmanager.h"
#include "common/treelandlogging.h"
#include "treelandconfig.hpp"

#include <winputdevice.h>
#include <wscoplistener.h>

#include <QSharedPointer>

#include <xkbcommon/xkbcommon.h>

#include <optional>

static QList<KeyboardStateWatcherV1 *> s_watchers;
static QHash<WSeat *, wlr_keyboard_modifiers> s_lastModifiers;


namespace {
bool isTreelandConfigInitialized(TreelandConfig *config)
{
    if (!config)
        return false;

#if TREELANDCONFIG_DCONFIG_FILE_VERSION_MINOR > 0
    return config->isInitializeSucceeded();
#else
    return config->isInitializeSucceed();
#endif
}

}

struct ModifierInfo {
    uint32_t flag;
    const char *modName;
};

static constexpr ModifierInfo s_modifierInfos[] = {
    { KeyboardStateWatcherV1::CapsLock, XKB_MOD_NAME_CAPS },
    { KeyboardStateWatcherV1::NumLock, XKB_MOD_NAME_NUM },
};

static std::optional<wlr_keyboard *> getSeatKeyboard(WSeat *seat)
{
    if (!seat || !seat->keyboardGroupKeyboard())
        return std::nullopt;

    auto *wlrDevice = seat->keyboardGroupKeyboard()->handle();
    if (wlrDevice->type != WLR_INPUT_DEVICE_KEYBOARD)
        return std::nullopt;
    auto *keyboard = wlr_keyboard_from_input_device(wlrDevice);
    if (!keyboard || !keyboard->keymap)
        return std::nullopt;

    return keyboard;
}

static bool modifiersEqual(const wlr_keyboard_modifiers &a, const wlr_keyboard_modifiers &b)
{
    return a.depressed == b.depressed
        && a.latched == b.latched
        && a.locked == b.locked
        && a.group == b.group;
}

static uint32_t lockStateBitfield(wlr_keyboard *keyboard, xkb_mod_mask_t locked)
{
    uint32_t state = 0;
    for (const auto &info : s_modifierInfos) {
        xkb_mod_index_t idx = xkb_keymap_mod_get_index(keyboard->keymap, info.modName);
        if (idx != XKB_MOD_INVALID && (locked & (xkb_mod_mask_t(1) << idx)))
            state |= info.flag;
    }

    return state;
}

static std::optional<uint32_t> getSeatLockState(WSeat *seat)
{
    const auto keyboard = getSeatKeyboard(seat);
    if (!keyboard)
        return std::nullopt;

    return lockStateBitfield(*keyboard, (*keyboard)->modifiers.locked);
}

static uint32_t changedLockStateBitfield(wlr_keyboard *keyboard,
                                         const wlr_keyboard_modifiers &previous,
                                         const wlr_keyboard_modifiers &current)
{
    return lockStateBitfield(keyboard, previous.locked ^ current.locked);
}

struct KeyboardConnection {
    WScopedListener modifiersListener;
    QPointer<WSeat> seat = nullptr;
};

class TreelandKeyboardStateNotifyManagerInterfaceV1Private
    : public QtWaylandServer::treeland_keyboard_state_notify_manager_v1
{
public:
    explicit TreelandKeyboardStateNotifyManagerInterfaceV1Private(TreelandKeyboardStateNotifyManagerInterfaceV1 *_q);

    TreelandKeyboardStateNotifyManagerInterfaceV1 *q = nullptr;

    void setupKeyboardConnections();
    void onModifiersEvent(WSeat *seat);
    wl_global *global() const;

protected:
    void bind_resource(Resource *resource) override;
    void destroy(Resource *resource) override;
    void get_keyboard_state_watcher(Resource *resource,
                                     uint32_t id,
                                     struct ::wl_resource *seat) override;

private:
    void handleSeatAdded(WSeat *seat);
    void handleSeatDestroy(WSeat *seat);
    void connectKeyboardGroup(WSeat *seat, WInputDevice *keyboardDevice);

    QList<QSharedPointer<KeyboardConnection>> m_keyboardConnections;
    bool m_keyboardConnectionsSetup = false;
};

TreelandKeyboardStateNotifyManagerInterfaceV1Private::TreelandKeyboardStateNotifyManagerInterfaceV1Private(
    TreelandKeyboardStateNotifyManagerInterfaceV1 *_q)
    : q(_q)
{
}

wl_global *TreelandKeyboardStateNotifyManagerInterfaceV1Private::global() const
{
    return m_global;
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::bind_resource([[maybe_unused]] Resource *resource)
{
    setupKeyboardConnections();
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::connectKeyboardGroup(WSeat *seat, WInputDevice *keyboardDevice)
{
    for (const auto &conn : std::as_const(m_keyboardConnections)) {
        if (conn->seat == seat)
            return;
    }

    if (const auto keyboard = getSeatKeyboard(seat))
        s_lastModifiers[seat] = (*keyboard)->modifiers;

    auto conn = QSharedPointer<KeyboardConnection>::create();
    conn->seat = seat;
    auto *wlrDevice = keyboardDevice->handle();
    if (wlrDevice->type == WLR_INPUT_DEVICE_KEYBOARD) {
        auto *keyboard = wlr_keyboard_from_input_device(wlrDevice);
        conn->modifiersListener.init(&keyboard->events.modifiers, q, [this, seat] (void *) {
            onModifiersEvent(seat);
        });
    }
    m_keyboardConnections.push_back(std::move(conn));
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::setupKeyboardConnections()
{
    if (m_keyboardConnectionsSetup)
        return;
    m_keyboardConnectionsSetup = true;

    auto *helper = Helper::instance();
    const auto seats = helper->seatManager()->seats();
    for (auto *seat : seats) {
        handleSeatAdded(seat);
    }

    QObject::connect(helper->seatManager(), &SeatsManager::seatAdded, q, [this](WSeat *seat) {
        handleSeatAdded(seat);
    });
    QObject::connect(helper->seatManager(), &SeatsManager::seatRemoved, q, [this](WSeat *seat) {
        handleSeatDestroy(seat);
    });
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::onModifiersEvent(WSeat *seat)
{
    const auto keyboard = getSeatKeyboard(seat);
    if (!keyboard)
        return;

    const auto currentModifiers = (*keyboard)->modifiers;
    const auto previousModifiers = s_lastModifiers.value(seat, currentModifiers);
    s_lastModifiers[seat] = currentModifiers;

    if (modifiersEqual(previousModifiers, currentModifiers))
        return;

    const uint32_t changedLocks = changedLockStateBitfield(*keyboard, previousModifiers, currentModifiers);
    if (changedLocks == 0)
        return;

    const uint32_t currentLocks = lockStateBitfield(*keyboard, currentModifiers.locked);

    constexpr uint32_t NumLockMask = KeyboardStateWatcherV1::NumLock;

    if (changedLocks & NumLockMask) {
        const bool isLocked = currentLocks & NumLockMask;

        if (isTreelandConfigInitialized(Helper::instance()->globalConfig())) {
            Helper::instance()->globalConfig()->setKeyboardNumLock(isLocked);
        }
    }

    if (s_watchers.isEmpty())
        return;

    for (auto *watcher : std::as_const(s_watchers)) {
        if (watcher->seat() && watcher->wSeat() != seat)
            continue;

        const uint32_t changed = changedLocks & uint32_t(watcher->modifiers());
        if (changed == 0)
            continue;

        for (int i = 0; i < 2; ++i) {
            uint32_t mod = (1u << i);
            if (!(changed & mod))
                continue;

            bool isLocked = (currentLocks & mod) != 0;
            if (isLocked && !(watcher->flags() & KeyboardStateWatcherV1::WatchLocked))
                continue;
            if (!isLocked && !(watcher->flags() & KeyboardStateWatcherV1::WatchUnlocked))
                continue;

            watcher->sendStateChanged(mod, isLocked ? KeyboardStateWatcherV1::Locked : KeyboardStateWatcherV1::Unlocked);
        }
    }
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::handleSeatDestroy(WSeat *seat)
{
    m_keyboardConnections.removeIf([seat](const QSharedPointer<KeyboardConnection> &conn) {
        return conn->seat == seat;
    });
    s_lastModifiers.remove(seat);
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::destroy(Resource *resource)
{
    wl_resource_destroy(resource->handle);
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::get_keyboard_state_watcher(Resource *resource,
                                                                                       uint32_t id,
                                                                                       struct ::wl_resource *seat)
{
    wl_resource *watcherResource = wl_resource_create(resource->client(),
                                                      &treeland_keyboard_state_watcher_v1_interface,
                                                      resource->version(),
                                                      id);
    if (!watcherResource) {
        wl_client_post_no_memory(resource->client());
        return;
    }

    auto *watcher = new KeyboardStateWatcherV1(watcherResource, seat);
    s_watchers.append(watcher);

    QObject::connect(watcher, &QObject::destroyed, [watcher]() {
        s_watchers.removeOne(watcher);
    });

    Q_EMIT q->keyboardStateWatcherCreated(watcher);
}

void TreelandKeyboardStateNotifyManagerInterfaceV1Private::handleSeatAdded(WSeat *seat)
{
    auto *keyboardevice = seat->keyboardGroupKeyboard();
    if (keyboardevice) {
        connectKeyboardGroup(seat, keyboardevice);
    }
}

class KeyboardStateWatcherV1Private : public QtWaylandServer::treeland_keyboard_state_watcher_v1
{
public:
    explicit KeyboardStateWatcherV1Private(KeyboardStateWatcherV1 *_q,
                                           wl_resource *resource,
                                           wl_resource *_seat);

    KeyboardStateWatcherV1 *q = nullptr;
    wl_resource *seat = nullptr;
    KeyboardStateWatcherV1::ModifierTypes modifiers;
    KeyboardStateWatcherV1::ModifierTypes pendingModifiers;
    KeyboardStateWatcherV1::WatchFlags flags;
    KeyboardStateWatcherV1::WatchFlags pendingFlags;

protected:
    void destroy_resource(Resource *resource) override;
    void destroy(Resource *resource) override;
    void set_modifiers(Resource *resource, uint32_t modifiers) override;
    void set_flags(Resource *resource, uint32_t flags) override;
    void apply(Resource *resource) override;
};

KeyboardStateWatcherV1Private::KeyboardStateWatcherV1Private(KeyboardStateWatcherV1 *_q,
                                                               wl_resource *resource,
                                                               wl_resource *_seat)
    : QtWaylandServer::treeland_keyboard_state_watcher_v1(resource)
    , q(_q)
    , seat(_seat)
{
}

void KeyboardStateWatcherV1Private::destroy_resource([[maybe_unused]] Resource *resource)
{
    delete q;
}

void KeyboardStateWatcherV1Private::destroy(Resource *resource)
{
    wl_resource_destroy(resource->handle);
}

void KeyboardStateWatcherV1Private::set_modifiers(Resource *resource, uint32_t modifiers)
{
    if (!treeland_keyboard_state_watcher_v1_modifier_is_valid(modifiers, resource->version())) {
        qCCritical(lcTlKeyboardNotify, "unknown modifiers");
        return;
    }

    pendingModifiers = KeyboardStateWatcherV1::ModifierTypes(modifiers);
}

void KeyboardStateWatcherV1Private::set_flags(Resource *resource, uint32_t flags)
{
    if (!treeland_keyboard_state_watcher_v1_watch_flag_is_valid(flags, resource->version())) {
        qCCritical(lcTlKeyboardNotify, "unknown watch flags");
        return;
    }

    pendingFlags = KeyboardStateWatcherV1::WatchFlags(flags);
}

void KeyboardStateWatcherV1Private::apply([[maybe_unused]] Resource *resource)
{
    if (pendingModifiers == modifiers && pendingFlags == flags)
        return;

    modifiers = pendingModifiers;
    flags = pendingFlags;

    if (modifiers == 0)
        return;

    auto *wSeat = q->wSeat();
    QList<WSeat *> seats;
    if (wSeat) {
        seats.append(wSeat);
    } else {
        auto *helper = Helper::instance();
        if (helper)
            seats = helper->seatManager()->seats();
    }

    for (auto *seat : std::as_const(seats)) {
        const auto currentLocks = getSeatLockState(seat);
        if (!currentLocks)
            continue;

        for (int i = 0; i < 2; ++i) {
            uint32_t mod = (1u << i);
            if (!(modifiers & mod))
                continue;

            bool isLocked = (*currentLocks & mod) != 0;
            if (isLocked && !(flags & KeyboardStateWatcherV1::WatchLocked))
                continue;
            if (!isLocked && !(flags & KeyboardStateWatcherV1::WatchUnlocked))
                continue;

            q->sendCurrentState(mod, isLocked ? KeyboardStateWatcherV1::Locked : KeyboardStateWatcherV1::Unlocked);
        }
    }
}

TreelandKeyboardStateNotifyManagerInterfaceV1::TreelandKeyboardStateNotifyManagerInterfaceV1(QObject *parent)
    : QObject(parent)
    , d(new TreelandKeyboardStateNotifyManagerInterfaceV1Private(this))
{
}

TreelandKeyboardStateNotifyManagerInterfaceV1::~TreelandKeyboardStateNotifyManagerInterfaceV1() = default;

void TreelandKeyboardStateNotifyManagerInterfaceV1::create(WServer *server)
{
    d->init(server->handle(), InterfaceVersion);
    d->setupKeyboardConnections();
}

void TreelandKeyboardStateNotifyManagerInterfaceV1::destroy([[maybe_unused]] WServer *server)
{
    d->globalRemove();
}

wl_global *TreelandKeyboardStateNotifyManagerInterfaceV1::global() const
{
    return d->global();
}

QByteArrayView TreelandKeyboardStateNotifyManagerInterfaceV1::interfaceName() const
{
    return d->interfaceName();
}

KeyboardStateWatcherV1::KeyboardStateWatcherV1(wl_resource *resource,
                                               wl_resource *seat,
                                               QObject *parent)
    : QObject(parent)
    , d(new KeyboardStateWatcherV1Private(this, resource, seat))
{
}

KeyboardStateWatcherV1::~KeyboardStateWatcherV1() = default;

KeyboardStateWatcherV1::ModifierTypes KeyboardStateWatcherV1::modifiers() const
{
    return d->modifiers;
}

KeyboardStateWatcherV1::WatchFlags KeyboardStateWatcherV1::flags() const
{
    return d->flags;
}

wl_resource *KeyboardStateWatcherV1::seat() const
{
    return d->seat;
}

WSeat *KeyboardStateWatcherV1::wSeat() const
{
    if (!d->seat)
        return nullptr;

    struct wlr_seat_client *seat_client =
        wlr_seat_client_from_resource(d->seat);
    Q_ASSERT_X(seat_client, __func__, "KeyboardStateWatcherV1 get wlr_seat_client failed.");
    return WSeat::fromHandle(seat_client->seat);
}

void KeyboardStateWatcherV1::sendStateChanged(uint32_t modifier, ModifierState state)
{
    d->send_state_changed(modifier, state);
}

void KeyboardStateWatcherV1::sendCurrentState(uint32_t modifier, ModifierState state)
{
    d->send_current_state(modifier, state);
}
