#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/drivers/i2s.h>
#include <zephyr/fs/fs.h>
#include <audio.h>
#include <string.h>

LOG_MODULE_REGISTER(audio, CONFIG_AUDIO_LOG_LEVEL);

#define SAMPLES_PER_BLOCK (CONFIG_AUDIO_BLOCK_SIZE / (CONFIG_AUDIO_WORD_WIDTH / 8) / 2) // Divide by 2 for stereo
#define BLOCK_DURATION_MS ((SAMPLES_PER_BLOCK * 1000U) / CONFIG_AUDIO_SAMPLE_RATE)      // Duration of audio in each block in milliseconds
#define MAX_WAIT_TIME_MS (3 * BLOCK_DURATION_MS)                                        // Maximum time to wait for the I2S peripheral to request the next block before we consider it stalled and reset it

/* Get the I2S device from the devicetree */
#define I2S_NODE DT_NODELABEL(i2s0)
static const struct device *i2s_dev = DEVICE_DT_GET(I2S_NODE);

/* Memory Slab for I2S DMA */
K_MEM_SLAB_DEFINE(audio_slab, CONFIG_AUDIO_BLOCK_SIZE, CONFIG_AUDIO_BLOCK_COUNT, 4);

/* Globals */
static volatile bool abort_playback = false;
static volatile uint8_t audio_volume = CONFIG_AUDIO_DEFAULT_VOLUME;

static void wait_for_i2s_drain(void)
{
    const uint32_t frames_per_block = CONFIG_AUDIO_BLOCK_SIZE / 4;
    const uint32_t block_ms = (frames_per_block * 1000U) / CONFIG_AUDIO_SAMPLE_RATE;
    const uint32_t max_wait_ms = (block_ms + 1U) * CONFIG_AUDIO_BLOCK_COUNT + 5U;
    int64_t deadline = k_uptime_get() + max_wait_ms;

    while (k_mem_slab_num_free_get(&audio_slab) < CONFIG_AUDIO_BLOCK_COUNT)
    {
        if (k_uptime_get() >= deadline)
        {
            LOG_WRN("Timeout waiting for I2S drain");
            break;
        }
        k_sleep(K_MSEC(1));
    }
}

/* Message Queue: transfers the file path to the thread */
K_MSGQ_DEFINE(audio_msgq, CONFIG_AUDIO_MAX_PATH_LEN, 10, 4);

/* Audio thread function */
void audio_thread_fn(void *p1, void *p2, void *p3)
{
    ARG_UNUSED(p1);
    ARG_UNUSED(p2);
    ARG_UNUSED(p3);

    char file_path[CONFIG_AUDIO_MAX_PATH_LEN];
    struct fs_file_t file;
    fs_file_t_init(&file);

    LOG_DBG("Audio thread started, priority %d", k_thread_priority_get(k_current_get()));
    while (1)
    {
        bool trigger_started = false;
        if (k_msgq_get(&audio_msgq, &file_path, K_FOREVER) == 0)
        {
            abort_playback = false;

            if (fs_open(&file, file_path, FS_O_READ) < 0)
            {
                LOG_ERR("thread: could not open %s", file_path);
                continue;
            }

            LOG_DBG("thread: preparing %s...", file_path);

            uint32_t queued_blocks = 0;

            while (!abort_playback)
            {
                void *mem_block;
                if (k_mem_slab_alloc(&audio_slab, &mem_block, K_MSEC(MAX_WAIT_TIME_MS)) < 0)
                {
                    LOG_ERR("audio: slab timeout (I2S stall? DMA failure?) - skipping sound and resetting I2S...");
                    i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_DROP);
                    audio_init();
                    break;
                }
                if (abort_playback)
                {
                    LOG_DBG("thread: playback aborted while waiting for memory block.");
                    k_mem_slab_free(&audio_slab, mem_block);
                    break;
                }

                int16_t *data_ptr = (int16_t *)mem_block;
                const uint32_t max_mono_samples = CONFIG_AUDIO_BLOCK_SIZE / 4;

                ssize_t bytes_read = fs_read(&file, data_ptr, max_mono_samples * sizeof(int16_t));

                if (bytes_read <= 0)
                {
                    k_mem_slab_free(&audio_slab, mem_block);
                    break;
                }

                uint32_t samples_read = bytes_read / sizeof(int16_t);

                // Padding with zeros if we read less than a full block of mono samples
                if (samples_read < max_mono_samples)
                {
                    memset(&data_ptr[samples_read], 0, (max_mono_samples - samples_read) * sizeof(int16_t));
                }

                uint32_t *stereo_dst = (uint32_t *)mem_block;
                for (int32_t i = max_mono_samples - 1; i >= 0; i--)
                {
                    int32_t scaled = (int32_t)data_ptr[i] * audio_volume;
                    int16_t sample = (int16_t)(scaled >> 8);

                    stereo_dst[i] = ((uint16_t)sample << 16) | (uint16_t)sample;
                }

                if (i2s_write(i2s_dev, mem_block, CONFIG_AUDIO_BLOCK_SIZE) < 0)
                {
                    k_mem_slab_free(&audio_slab, mem_block);
                    break;
                }

                queued_blocks++;

                // We start playback only when 2 blocks are in the DMA queue to avoid underruns
                if (!trigger_started && queued_blocks >= 2)
                {
                    if (i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_START) == 0)
                    {
                        trigger_started = true;
                        LOG_DBG("thread: playback started.");
                    }
                }

                if (samples_read < max_mono_samples)
                {
                    // Short sample: start with a single queued block so DRAIN can play it.
                    if (!trigger_started)
                    {
                        if (i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_START) == 0)
                        {
                            trigger_started = true;
                            LOG_DBG("thread: playback started (short sample).");
                        }
                    }
                    break;
                }
            }

            if (abort_playback)
            {
                i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_DROP);
                trigger_started = false;
                LOG_DBG("thread: playback aborted.");
            }
            else
            {
                if (k_msgq_num_used_get(&audio_msgq) > 0)
                {
                    LOG_DBG("thread: play request pending, not draining I2S to minimize latency...");
                }
                else
                {
                    LOG_DBG("thread: sample finished, waiting for I2S to drain...");
                    i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_DRAIN);
                    trigger_started = false;
                    wait_for_i2s_drain();
                    LOG_DBG("thread: playback finished.");
                }
            }
            fs_close(&file);
        }
    }
}

K_THREAD_DEFINE(audio_thread, CONFIG_AUDIO_STACK_SIZE, audio_thread_fn, NULL, NULL, NULL, CONFIG_AUDIO_THREAD_PRIORITY, 0, 0);

int audio_init(void)
{
    LOG_DBG("Initializing audio subsystem...");
    if (!device_is_ready(i2s_dev))
    {
        LOG_ERR("I2S device not ready");
        return -ENODEV;
    }

    /* Initial configuration of the I2S peripheral */
    struct i2s_config config = {
        .word_size = CONFIG_AUDIO_WORD_WIDTH,
        .channels = 2,
        .format = I2S_FMT_DATA_FORMAT_I2S,
        .options = I2S_OPT_BIT_CLK_MASTER | I2S_OPT_FRAME_CLK_MASTER,
        .frame_clk_freq = CONFIG_AUDIO_SAMPLE_RATE,
        .mem_slab = &audio_slab,
        .block_size = CONFIG_AUDIO_BLOCK_SIZE,
        .timeout = SYS_FOREVER_MS,
    };

    int ret = i2s_configure(i2s_dev, I2S_DIR_TX, &config);
    if (ret < 0)
    {
        LOG_ERR("Failed to configure I2S: %d", ret);
        return ret;
    }

    LOG_DBG("Audio subsystem initialized successfully");
    return 0;
}

int audio_play_file(const char *file)
{
    if (file == NULL)
    {
        LOG_ERR("audio_play_file: file path is NULL");
        return -EINVAL;
    }

    size_t len = strnlen(file, CONFIG_AUDIO_MAX_PATH_LEN);
    if (len >= CONFIG_AUDIO_MAX_PATH_LEN)
    {
        LOG_ERR("audio_play_file: file path too long: %s", file);
        return -ENAMETOOLONG;
    }

    char path_item[CONFIG_AUDIO_MAX_PATH_LEN];
    memcpy(path_item, file, len + 1);

    if (k_msgq_put(&audio_msgq, path_item, K_NO_WAIT) < 0)
    {
        LOG_ERR("audio_play_file: message queue full");
        return -EAGAIN;
    }

    LOG_DBG("Queued file for playback: %s", file);
    return 0;
}

int audio_play_sound(const char *file)
{
    if (file == NULL)
    {
        LOG_ERR("audio_play_sound: file name is NULL");
        return -EINVAL;
    }

    size_t len = strnlen(file, CONFIG_AUDIO_MAX_PATH_LEN) + strlen(CONFIG_FS_MGMT_MOUNT_POINT) + strlen(CONFIG_AUDIO_SAMPLE_FOLDER) + 2;
    if (len >= CONFIG_AUDIO_MAX_PATH_LEN)
    {
        LOG_ERR("audio_play_sound: file path too long: %s/%s/%s", CONFIG_FS_MGMT_MOUNT_POINT, CONFIG_AUDIO_SAMPLE_FOLDER, file);
        return -ENAMETOOLONG;
    }
    char path[CONFIG_AUDIO_MAX_PATH_LEN];

    snprintf(path, sizeof(path), "%s/%s/%s",
             CONFIG_FS_MGMT_MOUNT_POINT,
             CONFIG_AUDIO_SAMPLE_FOLDER,
             file);

    return audio_play_file(path);
}

void audio_stop(void)
{
    abort_playback = true;
    k_msgq_purge(&audio_msgq);
    i2s_trigger(i2s_dev, I2S_DIR_TX, I2S_TRIGGER_DROP);
    LOG_DBG("Playback stop requested, message queue purged");
}