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

#include <zmk/split/common.h>
#include <init.h>

#include <sys/crc.h>
#include <device.h>
#include <drivers/uart.h>

#include <logging/log.h>

LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);

#include <zmk/event_manager.h>
#include <zmk/events/position_state_changed.h>
#include <zmk/matrix.h>

#if !DT_HAS_CHOSEN(zmk_split_serial)
    #error "No zmk-split-serial node is chosen"
#endif

#define UART_NODE1 DT_CHOSEN(zmk_split_serial)
const struct device *serial_dev = DEVICE_DT_GET(UART_NODE1);
static int uart_ready = 0;

static void split_serial_receive_thread(void *unused, void *unused1, void *unused2);

K_MEM_SLAB_DEFINE(split_memory_slab, sizeof(split_data_t), \
                  CONFIG_ZMK_SPLIT_SERIAL_THREAD_QUEUE_SIZE, 4);

K_MSGQ_DEFINE(peripheral_event_msgq, sizeof(struct zmk_position_state_changed), \
              CONFIG_ZMK_SPLIT_SERIAL_THREAD_QUEUE_SIZE, 4);

K_THREAD_DEFINE(split_central, CONFIG_ZMK_SPLIT_SERIAL_THREAD_STACK_SIZE,
                split_serial_receive_thread, NULL, NULL, NULL,
                K_PRIO_PREEMPT(CONFIG_ZMK_SPLIT_SERIAL_THREAD_PRIORITY), 0, 0);


static void peripheral_event_work_callback(struct k_work *work) {
    struct zmk_position_state_changed ev;
    while (k_msgq_get(&peripheral_event_msgq, &ev, K_NO_WAIT) == 0) {
        LOG_DBG("Trigger key position state change for %d", ev.position);
        ZMK_EVENT_RAISE(new_zmk_position_state_changed(ev));
    }
}

K_WORK_DEFINE(peripheral_event_work, peripheral_event_work_callback);

static uint8_t split_central_notify_func(const void *data, uint16_t length) {
    static uint8_t position_state[SPLIT_DATA_LEN];
    uint8_t changed_positions[SPLIT_DATA_LEN];
    const split_data_t *split_data = data;
    uint16_t crc;

    LOG_DBG("[NOTIFICATION] data %p type:%u CRC:%u", data, split_data->type, split_data->crc);

    crc = crc16_ansi(split_data->data, sizeof(split_data->data));
    if (crc != split_data->crc) {
        LOG_WRN("CRC mismatch (%x:%x), skipping data", crc, split_data->crc);
        return 0;
    }


    for (int i = 0; i < SPLIT_DATA_LEN; i++) {
        changed_positions[i] = split_data->data[i] ^ position_state[i];
        position_state[i] = split_data->data[i];
    }

    for (int i = 0; i < SPLIT_DATA_LEN; i++) {
        for (int j = 0; j < 8; j++) {
            if (changed_positions[i] & BIT(j)) {
                uint32_t position = (i * 8) + j;
                bool pressed = position_state[i] & BIT(j);
                struct zmk_position_state_changed ev = {
                    .position = position, .state = pressed, .timestamp = k_uptime_get()};

                if (position > ZMK_KEYMAP_LEN) {
                    LOG_WRN("Invalid position: %u", position);
                    continue;
                }

                k_msgq_put(&peripheral_event_msgq, &ev, K_NO_WAIT);
                k_work_submit(&peripheral_event_work);
            }
        }
    }

    return 0;
}

static char *alloc_position_state_buffer() {
    char *block_ptr = NULL;
    if (k_mem_slab_alloc(&split_memory_slab, (void **)&block_ptr, K_NO_WAIT) == 0) {
        memset(block_ptr, 0, SPLIT_DATA_LEN);
    } else {
        LOG_WRN("Memory allocation time-out");
    }
    return block_ptr;
}


static void free_position_state_buffer(char *block_ptr) {
    k_mem_slab_free(&split_memory_slab, (void **)&block_ptr);
}


static void uart_callback(const struct device *dev, struct uart_event *evt, void *user_data) {
    char *buf = NULL;

    switch (evt->type) {

        case UART_RX_STOPPED:
            LOG_DBG("UART device:%s rx stopped", serial_dev->name);
            break;

        case UART_RX_BUF_REQUEST:
            LOG_DBG("UART device:%s rx extra buf req", serial_dev->name);
            buf = alloc_position_state_buffer();
            if (NULL != buf) {
                int ret = uart_rx_buf_rsp(serial_dev, buf, sizeof(split_data_t));
                if (0 != ret) {
                    LOG_WRN("UART device:%s rx extra buf req add failed: %d", serial_dev->name, ret);
                    free_position_state_buffer(buf);
                }
            }
            break;

        case UART_RX_RDY:
            LOG_DBG("UART device:%s rx buf ready", serial_dev->name);
            break;

        case UART_RX_BUF_RELEASED:
            LOG_DBG("UART device:%s rx buf released", serial_dev->name);
            split_central_notify_func(evt->data.rx_buf.buf, sizeof(split_data_t));
            free_position_state_buffer(evt->data.rx_buf.buf);
            break;

        default:
            LOG_DBG("UART device:%s unhandled event: %u", serial_dev->name, evt->type);
            break;
    };
    return;
}


static void split_serial_receive_thread(void *unused, void *unused1, void *unused2) {
    if (!device_is_ready(serial_dev)) {
        LOG_WRN("UART device:%s not ready", serial_dev->name);
        return;
    }

    int ret = uart_callback_set(serial_dev, uart_callback, NULL);
    if (ret == -ENOTSUP || ret == -ENOSYS) {
        LOG_WRN("UART device:%s ASYNC not supported", serial_dev->name);
        return;
    }

    uart_ready = 1;
    LOG_DBG("UART device:%s ready", serial_dev->name);

    while(true) {
        char *buf = alloc_position_state_buffer();
        if (NULL == buf) {
            k_msleep(100);
            continue;
        }

        int ret = uart_rx_enable(serial_dev, buf, sizeof(split_data_t), SYS_FOREVER_MS);
        if (ret == -ENOTSUP) {
            LOG_WRN("UART device:%s not supporting DMA", serial_dev->name);
            free_position_state_buffer(buf);
            return;
        }
        if (ret != 0 && ret != -EBUSY) {
            LOG_WRN("UART device:%s RX error:%d", serial_dev->name, ret);
            free_position_state_buffer(buf);
            continue;
        }
    };
}