#include <zephyr/device.h>
#include <zephyr/init.h>
#include <zephyr/drivers/gpio.h>
#include <zephyr/drivers/i2c.h>
#include <zephyr/drivers/i2s.h>
#include <zephyr/audio/codec.h>
#include <zephyr/logging/log.h>
#include <zephyr/sys/iterable_sections.h>

#include "wavheader.h"
#include "audio.h"

LOG_MODULE_REGISTER(audio, CONFIG_AUDIO_CODEC_LOG_LEVEL);

static const struct device *i2s_dev;
static int16_t *tx_block;
static struct k_work audio_work;
static int16_t *asset;
static uint32_t file_pos = 0;
static uint32_t asset_size = 0;

#define SAMPLE_NO 64

#define NUM_BLOCKS 20
#define BLOCK_SIZE (2 * sizeof(uint16_t) * SAMPLE_NO)

#ifdef CONFIG_NOCACHE_MEMORY
#define MEM_SLAB_CACHE_ATTR __nocache
#else
#define MEM_SLAB_CACHE_ATTR
#endif /* CONFIG_NOCACHE_MEMORY */

static char MEM_SLAB_CACHE_ATTR __aligned(WB_UP(32))
    _k_mem_slab_buf_tx_0_mem_slab[(NUM_BLOCKS)*WB_UP(BLOCK_SIZE)];

static STRUCT_SECTION_ITERABLE(k_mem_slab, tx_0_mem_slab) = Z_MEM_SLAB_INITIALIZER(
    tx_0_mem_slab, _k_mem_slab_buf_tx_0_mem_slab, WB_UP(BLOCK_SIZE), NUM_BLOCKS);

static void i2s_timer_handler(struct k_timer *timer_id);
K_TIMER_DEFINE(i2s_timer, i2s_timer_handler, NULL);

static void fill_buf(int att) {
    for (int i = 0; i < SAMPLE_NO; i += 1) {
        if (file_pos + i > asset_size) {
            tx_block[2 * i] = 0;
            tx_block[2 * i + 1] = 0;
            if (k_timer_remaining_get(&i2s_timer) > 0) {
                k_timer_stop(&i2s_timer);
            }
        } else {
            tx_block[2 * i] = (asset[file_pos + i] >> att);
            tx_block[2 * i + 1] = (asset[file_pos + i] >> att);
        }
    }
    file_pos += SAMPLE_NO;
}

static void audio_work_handler(struct k_work *work) {
    int ret;
    ret = k_mem_slab_alloc(&tx_0_mem_slab, (void *)&tx_block, K_FOREVER);
    if (ret != 0) {
        LOG_ERR("k_mem_slab_alloc: error %08x", ret);
        return;
    }

    fill_buf(0);

    ret = i2s_write(i2s_dev, tx_block, BLOCK_SIZE);
    if (ret != 0) {
        LOG_ERR("i2s_write: error %08x", ret);
        k_mem_slab_free(&tx_0_mem_slab, (void *)tx_block);
    }
}

static void i2s_timer_handler(struct k_timer *timer_id) { k_work_submit(&audio_work); }

static WavFileHeader_t *audio_file_init(const int16_t *asset_in, uint32_t asset_size_in) {
    asset_size = asset_size_in;

    LOG_INF("sizeof(WavFileHeader_t): %d", sizeof(WavFileHeader_t));
    LOG_INF("asset loc: %p size: %d", (void *)asset_in, asset_size);

    WavFileHeader_t *wavheader = (WavFileHeader_t *)asset_in;
    LOG_INF("PcmFlags: %d Channels: %d", wavheader->PcmFlags, wavheader->Channels);
    LOG_INF("BitDepth: %d SampleRate: %d", wavheader->BitDepth, wavheader->SampleRate);
    // LOG_INF("FileFormat: %s BlockAlign: %d", wavheader->FileFormat, wavheader->BlockAlign);
    LOG_INF("\n");

    asset = asset_in + sizeof(WavFileHeader_t) / 2;
    return wavheader;
}

static void audio_i2s_init(const struct device *dev, WavFileHeader_t *wavheader) {
    i2s_dev = dev;
    struct i2s_config i2s_cfg;
    int ret;

    i2s_cfg.word_size = 16U;
    i2s_cfg.channels = 2U;
    i2s_cfg.format = I2S_FMT_DATA_FORMAT_I2S;
    i2s_cfg.frame_clk_freq = wavheader->SampleRate;
    i2s_cfg.block_size = BLOCK_SIZE;
    i2s_cfg.timeout = 2000;
    i2s_cfg.options = I2S_OPT_FRAME_CLK_MASTER | I2S_OPT_BIT_CLK_MASTER;
    i2s_cfg.mem_slab = &tx_0_mem_slab;
    ret = i2s_configure(i2s_dev, I2S_DIR_TX, &i2s_cfg);
    if (ret < 0) {
        LOG_ERR("Failed to configure I2S stream\n");
        return ret;
    }
}

int audio_play(const struct device *dev, const int16_t *asset_in, uint32_t asset_size_in) {

    WavFileHeader_t *wavheader = audio_file_init(asset_in, asset_size_in);
    audio_i2s_init(dev, wavheader);
    file_pos = 0;

    k_work_init(&audio_work, audio_work_handler);

    k_timer_start(&i2s_timer, K_USEC(5), K_USEC(5));

    audio_work_handler(NULL); // prime the i2s buffer

    /* Trigger the I2S transmission */
    int ret = i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_START);
    if (ret < 0) {
        LOG_ERR("Could not trigger I2S tx\n");
        return ret;
    }

    // TODO replace with DONE flag from timer handler and ditch loop
    // had issues prototyping this, stashed away for now.
    uint32_t cnt = 0x800;
    while (1) {

        k_sleep(K_MSEC(1));
        if (--cnt == 0) {
            if (k_timer_remaining_get(&i2s_timer) > 0) {
                k_timer_stop(&i2s_timer);
            }
            break;
        }
    }

    /* Drain TX queue */
    // ret = i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_DRAIN);
    // if (ret < 0) {
    //     LOG_ERR("Could not trigger I2S tx\n");
    //     return ret;
    // }
    LOG_DBG("All I2S blocks written\n");

    return 0;
}