#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/jiffies.h>
#include <linux/delay.h>
#include <linux/spinlock.h>
#include <linux/wait.h>
#include <linux/completion.h>
#include <linux/kthread.h>
#include <linux/interrupt.h>
#include <linux/device.h>
#include <linux/cdev.h>
#include <linux/spi/spi.h>
#include <linux/gpio.h>

#include <linux/of_device.h>
#include <linux/uaccess.h>

#include "msgq.h"
#include "esp32.h"

#define wait_event_interruptible_lock(wq_head, condition, lock)             \
({                                                                          \
    int __ret = 0;                                                          \
    if (!(condition))                                                       \
        __ret = ___wait_event(wq_head, condition, TASK_INTERRUPTIBLE, 0, 0, \
                                spin_unlock(&lock);                         \
                                schedule();                                 \
                                spin_lock(&lock));                          \
    __ret;                                                                  \
})

/* has to be power of 2 */
#define MSGQ_SIZE_SND       128
#define MSGQ_SIZE_RCV       128

#define ESP32_SIZE_HDR      3
#define ESP32_SIZE_BUF      (ESP32_MTU + 4)

static int esp32_open(struct inode *inode, struct file *file);
static int esp32_release(struct inode *inode, struct file *file);
static ssize_t esp32_read(struct file *file, char __user *buf, size_t len, loff_t *off);
static ssize_t esp32_write(struct file *file, const char __user *buf, size_t len, loff_t *off);

static int esp32_probe(struct spi_device *spi);
static void esp32_remove(struct spi_device *spi);

static DECLARE_WAIT_QUEUE_HEAD(xchg_wq);
static DECLARE_WAIT_QUEUE_HEAD(cts_wq);

static MSGQueue snd_msgq;
static MSGQueue rcv_msgq[ESP32_MAX_MTYPE];

static struct spi_device *esp32_spi_dev;
static struct gpio_desc *gpio_rts, *gpio_cts;
static int irq_rts, irq_cts;
static int cts_val;
static struct task_struct *kthd;

static dev_t dev;
static struct class *dev_class;
static struct cdev dev_cdev;
static struct file_operations dev_fops = {
    .owner      = THIS_MODULE,
    .open       = esp32_open,
    .release    = esp32_release,
    .read       = esp32_read,
    .write      = esp32_write,
};

static const struct of_device_id esp32_of_match[] = {
    {
        .compatible = "mikrophone,esp32",
    },
    { /* sentinel */ }
};
MODULE_DEVICE_TABLE(of, esp32_of_match);

static const struct spi_device_id esp32_spi_ids[] = {
    { "esp32", 0 },
    { /* sentinel */ }
};
MODULE_DEVICE_TABLE(spi, esp32_spi_ids);

static struct spi_driver esp32_spi_driver = {
    .probe = esp32_probe,
    .remove = esp32_remove,
    .driver = {
        .name = "esp32",
        .of_match_table = of_match_ptr(esp32_of_match),
    },
    .id_table = esp32_spi_ids,
};

static irqreturn_t rts_irq_handler(int irq, void *dev_id) {
    wake_up(&xchg_wq);

    return IRQ_HANDLED;
}

static irqreturn_t cts_irq_handler(int irq, void *dev_id) {
    cts_val = 0;
    wake_up(&cts_wq);

    return IRQ_HANDLED;
}

/* esp32 dma workaraund */
static inline unsigned int esp32_fix_len(unsigned int len) {
    if (len < 8) {
        return 8;
    } else if (len % 4 != 0) {
        return (len / 4 + 1) * 4;
    }
    return len;
}

static inline int wait4cts(void) {
    return wait_event_interruptible(cts_wq, !cts_val);
}

static inline int spi_xchg(struct spi_transfer *spi_tr) {
    cts_val = 1;
    return spi_sync_transfer(esp32_spi_dev, spi_tr, 1);
}

static int xchg_thd(void *arg) {
    unsigned char *buffer, *_buffer, *buf_rx;
    unsigned char mtype, mtype_flags;
    int rv, wake_sndq, wake_rcvq, bfree;
    struct spi_transfer spi_tr;
    uint16_t len_tx, len_rx;

    _buffer = NULL;
    memset(&spi_tr, 0, sizeof(spi_tr));

    buf_rx = esp32_alloc();
    if (buf_rx == NULL) return -ENOMEM;

    while (!kthread_should_stop()) {
        /* preallocate buffer outside of spin lock: GFP_KERNEL allocator can sleep */
        if (_buffer == NULL) {
            _buffer = esp32_alloc();
            if (_buffer == NULL) return -ENOMEM;
            _buffer[0] = 0;
            _buffer[1] = 0;
            _buffer[2] = 0;
        }

        rv = 0;
        wake_sndq = 0;
        wake_rcvq = 0;
        bfree = 1;

        spin_lock(&snd_msgq.lock);
        do {
            buffer = msgq_pop(&snd_msgq);
            if (buffer && (msgq_len(&snd_msgq) == (msgq_size(&snd_msgq) - 1))) {
                wake_sndq = 1;
            }
            if ((buffer == NULL) && (gpiod_get_value(gpio_rts) == 0)) {
                buffer = _buffer;
                _buffer = NULL;
            }
            if (buffer == NULL) {
                rv = wait_event_interruptible_lock(xchg_wq, (msgq_len(&snd_msgq) != 0) || (gpiod_get_value(gpio_rts) == 0), snd_msgq.lock);
                if (rv) break;
            }
        } while (buffer == NULL);
        spin_unlock(&snd_msgq.lock);

        /* kthread_should_stop() might be true */
        if (rv) continue;

        if (wake_sndq) wake_up(&snd_msgq.wq);

        len_tx  = (uint16_t)buffer[1] << 8;
        len_tx |= (uint16_t)buffer[2] & 0xFF;
        if (len_tx > ESP32_MTU) goto xchg_thd_fin;

        spi_tr.tx_buf = buffer;
        spi_tr.rx_buf = buf_rx;
        spi_tr.len = len_tx + ESP32_SIZE_HDR;
        spi_tr.len = esp32_fix_len(spi_tr.len);

        /* wait for cts */
        rv = wait4cts();
        if (rv) goto xchg_thd_fin;

        /* nothing to send/receive */
        if ((buffer[0] == 0) && gpiod_get_value(gpio_rts)) goto xchg_thd_fin;

        rv = spi_xchg(&spi_tr);
        if (rv) goto xchg_thd_fin;

        mtype = buf_rx[0] & ESP32_MTYPE_MASK;
        mtype_flags = buf_rx[0] & ~ESP32_MTYPE_MASK;
        if (mtype >= ESP32_MAX_MTYPE) goto xchg_thd_fin;

        if (mtype) {
            memcpy(buffer, spi_tr.rx_buf, spi_tr.len);

            len_rx  = (uint16_t)buffer[1] << 8;
            len_rx |= (uint16_t)buffer[2] & 0xFF;
            if (len_rx > ESP32_MTU) goto xchg_thd_fin;

            if (len_rx + ESP32_SIZE_HDR > spi_tr.len) {
                spi_tr.tx_buf = NULL;
                spi_tr.rx_buf = buffer + spi_tr.len;
                spi_tr.len = len_rx + ESP32_SIZE_HDR - spi_tr.len;
                spi_tr.len = esp32_fix_len(spi_tr.len);

                /* wait for cts, again */
                rv = wait4cts();
                if (rv) goto xchg_thd_fin;

                rv = spi_xchg(&spi_tr);
                if (rv) goto xchg_thd_fin;
            }

            spin_lock(&rcv_msgq[mtype].lock);
            rv = msgq_push(&rcv_msgq[mtype], buffer);
            if (!rv && (msgq_len(&rcv_msgq[mtype]) == 1)) {
                wake_rcvq = 1;
            }
            spin_unlock(&rcv_msgq[mtype].lock);

            if (!rv) bfree = 0;
            if (wake_rcvq) wake_up(&rcv_msgq[mtype].wq);
        }

xchg_thd_fin:
        if (bfree) {
            if (_buffer == NULL) {
                _buffer = buffer;
                _buffer[0] = 0;
                _buffer[1] = 0;
                _buffer[2] = 0;
            } else {
                esp32_free(buffer);
            }
        }
    }

    esp32_free(_buffer);
    esp32_free(buf_rx);
    return 0;
}

static int esp32_open(struct inode *inode, struct file *file) {
    unsigned int major, minor;

    major = MAJOR(dev);
    minor = MINOR(dev);
    if (imajor(inode) != major) return -EINVAL;
    if ((iminor(inode) < minor) || (iminor(inode) >= minor + ESP32_MAX_MTYPE)) return -EINVAL;
    file->private_data = (void *)(iminor(inode) - minor);

    return 0;
}

static int esp32_release(struct inode *inode, struct file *file) {
    return 0;
}

static ssize_t esp32_read(struct file *file, char __user *buf, size_t len, loff_t *off) {
    unsigned char mtype;
    unsigned char *buffer;
    uint16_t msg_len;
    unsigned long _rv;
    int block, rv = 0;

    block = !(file->f_flags & O_NONBLOCK);
    mtype = (unsigned char)(file->private_data);
    rv = esp32_receive(mtype, block, &buffer);
    if (rv) goto esp32_read_fin;

    msg_len  = (uint16_t)buffer[1] << 8;
    msg_len |= (uint16_t)buffer[2] & 0xFF;
    if (len < msg_len) rv = -EINVAL;
    if (rv) goto esp32_read_fin;

    _rv = copy_to_user(buf, buffer + ESP32_SIZE_HDR, msg_len);
    if (_rv) rv = -EFAULT;

esp32_read_fin:
    if (buffer) esp32_free(buffer);
    if (rv) return rv;
    return msg_len;
}

static ssize_t esp32_write(struct file *file, const char __user *buf, size_t len, loff_t *off) {
    unsigned char mtype;
    unsigned char *buffer;
    unsigned long _rv;
    int block, rv = 0;

    if (len > ESP32_MTU) return -EINVAL;

    buffer = esp32_alloc();
    if (buffer == NULL) return -ENOMEM;

    _rv = copy_from_user(buffer + ESP32_SIZE_HDR, buf, len);
    if (_rv) rv = -EFAULT;
    if (rv) {
        esp32_free(buffer);
        return rv;
    }

    block = !(file->f_flags & O_NONBLOCK);
    mtype = (unsigned char)(file->private_data);
    buffer[0] = mtype;
    buffer[1] = len >> 8;
    buffer[2] = len & 0xFF;

    rv = esp32_send(buffer, block);
    if (rv) return rv;

    return len;
}

static int esp32_probe(struct spi_device *spi) {
    dev_t _dev;
    struct device *dev_ptr;
    const char *fmt;
    int i, rv, major, minor;
    const struct of_device_id *match;

    match = of_match_device(of_match_ptr(esp32_of_match), &spi->dev);
    if (match == NULL) return -ENODEV;

    esp32_spi_dev = spi;

    gpio_rts = gpiod_get(&spi->dev, "rts", GPIOD_IN);
    if (IS_ERR(gpio_rts)) goto err_gpiod_get_rts;

    irq_rts = gpiod_to_irq(gpio_rts);
    if (irq_rts < 0) goto err_gpiod_get_rts;

    gpio_cts = gpiod_get(&spi->dev, "cts", GPIOD_IN);
    if (IS_ERR(gpio_cts) || gpiod_cansleep(gpio_rts)) goto err_gpiod_get_cts;

    irq_cts = gpiod_to_irq(gpio_cts);
    if (irq_cts < 0) goto err_gpiod_get_cts;

    rv = request_irq(irq_rts, rts_irq_handler, IRQF_TRIGGER_FALLING, "esp32_rts", NULL);
    if (rv < 0) goto err_request_irq_rts;

    rv = request_irq(irq_cts, cts_irq_handler, IRQF_TRIGGER_FALLING, "esp32_cts", NULL);
    if (rv < 0) goto err_request_irq_cts;

    kthd = kthread_run(xchg_thd, NULL, "esp32");
    if (IS_ERR(kthd)) goto err_kthread_run;

    rv = alloc_chrdev_region(&dev, 0, ESP32_MAX_MTYPE, "esp32");
    if (rv) goto err_alloc_chrdev;

    cdev_init(&dev_cdev, &dev_fops);

    rv = cdev_add(&dev_cdev, dev, ESP32_MAX_MTYPE);
    if (rv) goto err_cdev_add;

    dev_class = class_create("esp32");
    if (IS_ERR(dev_class)) goto err_class_create;

    major = MAJOR(dev);
    minor = MINOR(dev);
    for (i=0; i<ESP32_MAX_MTYPE; i++) {
        _dev = MKDEV(major, minor + i);
        fmt = i & ESP32_MTYPE_FLAG_BRIDGE ? "fe310.%d" : "esp32.%d";
        dev_ptr = device_create(dev_class, NULL, _dev, NULL, fmt, i & ~ESP32_MTYPE_FLAG_BRIDGE);
        if (IS_ERR(dev_ptr)) goto err_device_create;
    }

    return 0;

err_device_create:
    for (i--; i>=0; i--) {
        _dev = MKDEV(major, minor + i);
        device_destroy(dev_class, _dev);
    }
    class_destroy(dev_class);

err_class_create:
    cdev_del(&dev_cdev);

err_cdev_add:
    unregister_chrdev_region(dev, ESP32_MAX_MTYPE);

err_alloc_chrdev:
    kthread_stop(kthd);

err_kthread_run:
    free_irq(irq_cts, NULL);

err_request_irq_cts:
    free_irq(irq_rts, NULL);

err_request_irq_rts:
err_gpiod_get_cts:
    if (!IS_ERR(gpio_cts)) gpiod_put(gpio_rts);

err_gpiod_get_rts:
    if (!IS_ERR(gpio_rts)) gpiod_put(gpio_rts);
    return -1;
}

static void esp32_remove(struct spi_device *spi) {
    unsigned char *buffer;
    dev_t _dev;
    int i, major, minor;

    major = MAJOR(dev);
    minor = MINOR(dev);
    for (i=0; i<ESP32_MAX_MTYPE; i++) {
        _dev = MKDEV(major, minor + i);
        device_destroy(dev_class, _dev);
    }
    class_destroy(dev_class);
    cdev_del(&dev_cdev);
    unregister_chrdev_region(dev, ESP32_MAX_MTYPE);
    kthread_stop(kthd);
    free_irq(irq_cts, NULL);
    free_irq(irq_rts, NULL);
    gpiod_put(gpio_cts);
    gpiod_put(gpio_rts);

    /* flush remaining buffers */
    spin_lock(&snd_msgq.lock);
    while ((buffer = msgq_pop(&snd_msgq))) {
        esp32_free(buffer);
    }
    msgq_destroy(&snd_msgq);
    spin_unlock(&snd_msgq.lock);
    for (i=0; i<ESP32_MAX_MTYPE; i++) {
        spin_lock(&rcv_msgq[i].lock);
        while ((buffer = msgq_pop(&rcv_msgq[i]))) {
            esp32_free(buffer);
        }
        msgq_destroy(&rcv_msgq[i]);
        spin_unlock(&rcv_msgq[i].lock);
    }
}

static int __init esp32_init(void) {
    int i, rv;

    rv = msgq_create(&snd_msgq, MSGQ_SIZE_SND);
    if (rv) return rv;

    for (i=0; i<ESP32_MAX_MTYPE; i++) {
        rv = msgq_create(&rcv_msgq[i], MSGQ_SIZE_RCV);
        if (rv) goto esp32_init_err;
    }

    rv = spi_register_driver(&esp32_spi_driver);
    if (rv) goto esp32_init_err;

    return 0;

esp32_init_err:
    for (i--; i>=0; i--) {
        msgq_destroy(&rcv_msgq[i]);
    }
    msgq_destroy(&snd_msgq);
    return rv;
}

static void __exit esp32_exit(void) {
    spi_unregister_driver(&esp32_spi_driver);
}

module_init(esp32_init);
module_exit(esp32_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Uros Majstorovic <majstor@majstor.org>");
MODULE_DESCRIPTION("esp32 SPI driver");
MODULE_VERSION("1.0");

int esp32_send(unsigned char *buffer, int block) {
    int rv, wake_xchg;

    spin_lock(&snd_msgq.lock);
    do {
        wake_xchg = 0;
        rv = msgq_push(&snd_msgq, buffer);
        if (!rv && (msgq_len(&snd_msgq) == 1)) {
            wake_xchg = 1;
        }
        if (block && rv) {
            int _rv;

            _rv = wait_event_interruptible_lock(snd_msgq.wq, msgq_len(&snd_msgq) != msgq_size(&snd_msgq), snd_msgq.lock);
            if (_rv) {
                rv = _rv;
                break;
            }
        }
    } while (block && rv);
    spin_unlock(&snd_msgq.lock);

    if (wake_xchg) wake_up(&xchg_wq);
    if (rv) esp32_free(buffer);

    return rv;
}

int esp32_receive(unsigned char mtype, int block, unsigned char **buffer) {
    int rv;

    *buffer = NULL;
    if (mtype >= ESP32_MAX_MTYPE) return -EINVAL;

    rv = 0;
    spin_lock(&rcv_msgq[mtype].lock);
    *buffer = msgq_pop(&rcv_msgq[mtype]);
    while (block && (*buffer == NULL)) {
        rv = wait_event_interruptible_lock(rcv_msgq[mtype].wq, msgq_len(&rcv_msgq[mtype]) != 0, rcv_msgq[mtype].lock);
        if (rv) break;

        *buffer = msgq_pop(&rcv_msgq[mtype]);
    }
    spin_unlock(&rcv_msgq[mtype].lock);

    if (!rv && (*buffer == NULL)) rv = -EAGAIN;
    return rv;
}

unsigned char *esp32_alloc(void) {
    return kmalloc(ESP32_SIZE_BUF, GFP_KERNEL);
}

void esp32_free(unsigned char *buffer) {
    kfree(buffer);
}