/*
 * Copyright (c) 2021 The ZMK Contributors
 *
 * SPDX-License-Identifier: MIT
 */

#define DT_DRV_COMPAT zmk_conditional_layers

#include <stdint.h>
#include <zephyr/kernel.h>

#include <zephyr/devicetree.h>
#include <zephyr/logging/log.h>

#include <zmk/event_manager.h>
#include <zmk/keymap.h>
#include <zmk/events/layer_state_changed.h>

LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);

#if DT_HAS_COMPAT_STATUS_OKAY(DT_DRV_COMPAT)

static K_SEM_DEFINE(conditional_layer_sem, 1, 1);

// Conditional layer configuration that activates the specified then-layer when all if-layers are
// active. With two if-layers, this is referred to as "tri-layer", and is commonly used to activate
// a third "adjust" layer if and only if the "lower" and "raise" layers are both active.
struct conditional_layer_cfg {
    // A bitmask of each layer that must be pressed for this conditional layer config to activate.
    zmk_keymap_layers_state_t if_layers_state_mask;

    // The layer number that should be active while all layers in the if-layers mask are active.
    int8_t then_layer;
};

#define IF_LAYER_BIT(node_id, prop, idx) BIT(DT_PROP_BY_IDX(node_id, prop, idx)) |

// Evaluates to conditional_layer_cfg struct initializer.
#define CONDITIONAL_LAYER_DECL(n)                                                                  \
    {                                                                                              \
        .if_layers_state_mask = DT_FOREACH_PROP_ELEM(n, if_layers, IF_LAYER_BIT) 0,                \
        .then_layer = DT_PROP(n, then_layer),                                                      \
    },

// All conditional layer configurations in the keymap.
static const struct conditional_layer_cfg CONDITIONAL_LAYER_CFGS[] = {
    DT_INST_FOREACH_CHILD(0, CONDITIONAL_LAYER_DECL)};

static const int32_t NUM_CONDITIONAL_LAYER_CFGS =
    sizeof(CONDITIONAL_LAYER_CFGS) / sizeof(*CONDITIONAL_LAYER_CFGS);

// Ensures all layer updates are processed if one conditional layer activates another.
static bool conditional_layer_updates_needed;
// Tracks which layers have been locked by conditional layer activations.
static uint32_t layer_locked_by_conditional = 0;

static void conditional_layer_activate(int8_t layer, bool locking) {
    // This may trigger another event that could, in turn, activate additional then-layers. However,
    // the process will eventually terminate (at worst, when every layer is active).
    if (!zmk_keymap_layer_active(layer) || (locking && !zmk_keymap_layer_locked(layer))) {
        LOG_DBG("layer %d", layer);
        zmk_keymap_layer_activate(layer, locking);
    }
}

static void conditional_layer_deactivate(int8_t layer, bool locking) {
    // This may deactivate a then-layer that's already active via another mechanism (e.g., a
    // momentary layer behavior). However, the same problem arises when multiple keys with the same
    // &mo binding are held and then one is released, so it's probably not an issue in practice.
    if (zmk_keymap_layer_active(layer) && (!zmk_keymap_layer_locked(layer) || locking)) {
        LOG_DBG("layer %d", layer);
        zmk_keymap_layer_deactivate(layer, locking);
    }
}

static int layer_state_changed_listener(const zmk_event_t *ev) {
    conditional_layer_updates_needed = true;

    // Semaphore ensures we don't re-enter the loop in the middle of doing update, and
    // ensures that "waterfalling layer updates" are all processed to trigger subsequent
    // nested conditional layers properly.
    if (k_sem_take(&conditional_layer_sem, K_NO_WAIT) < 0) {
        return 0;
    }

    while (conditional_layer_updates_needed) {
        int8_t max_then_layer = -1;
        uint32_t then_layers = 0;
        uint32_t then_layer_state = 0;
        conditional_layer_updates_needed = false;

        // On layer state changes, examines each conditional layer config to determine if then-layer
        // in the config should activate based on the currently active set of if-layers.
        for (int i = 0; i < NUM_CONDITIONAL_LAYER_CFGS; i++) {
            const struct conditional_layer_cfg *cfg = CONDITIONAL_LAYER_CFGS + i;
            zmk_keymap_layers_state_t mask = cfg->if_layers_state_mask;
            WRITE_BIT(then_layers, cfg->then_layer, true);
            max_then_layer = MAX(max_then_layer, cfg->then_layer);

            // Activate then-layer if and only if all if-layers are already active. Note that we
            // reevaluate the current layer state for each config since activation of one layer can
            // also trigger activation of another.
            if ((zmk_keymap_layer_state() & mask) == mask) {
                WRITE_BIT(then_layer_state, cfg->then_layer, true);
            }
            // Same as above, but for the lock status
            if ((zmk_keymap_layer_locks() & mask) == mask) {
                WRITE_BIT(layer_locked_by_conditional, cfg->then_layer, true);
            }
        }

        for (uint8_t layer = 0; layer <= max_then_layer; layer++) {
            if ((BIT(layer) & then_layers) != 0U) {
                bool locking = (BIT(layer) & layer_locked_by_conditional) != 0U;
                if ((BIT(layer) & then_layer_state) != 0U) {
                    conditional_layer_activate(layer, locking);
                } else {
                    conditional_layer_deactivate(layer, locking);
                    WRITE_BIT(layer_locked_by_conditional, layer, false);
                }
            }
        }
    }

    k_sem_give(&conditional_layer_sem);
    return 0;
}

ZMK_LISTENER(conditional_layer, layer_state_changed_listener);
ZMK_SUBSCRIPTION(conditional_layer, zmk_layer_state_changed);

#endif