/*
 * Espressif Modified MIT License
 *
 * Copyright (c) 2025 Espressif Systems (Shanghai) Co., LTD
 *
 * Permission is hereby granted for use **exclusively** with Espressif Systems products.
 * This includes the right to use, copy, modify, merge, publish, distribute, and sublicense
 * the Software, subject to the following conditions:
 *
 * 1. This Software **must be used in conjunction with Espressif Systems products**.
 * 2. The above copyright notice and this permission notice shall be included in all copies
 *    or substantial portions of the Software.
 * 3. Redistribution of the Software in source or binary form **for use with non-Espressif products**
 *    is strictly prohibited.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
 * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
 * FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 *
 * SPDX-License-Identifier: MIT-ESPRESSIF
 */

#include <stdio.h>
#include "esp_log.h"
#include "esp_check.h"
#include "audio_pipeline.h"
#include "raw_stream.h"
#include "filter_resample.h"
#include "algorithm_stream.h"
#include "raw_stream.h"
#include "tone_stream.h"
#include "i2s_stream.h"
#include "mp3_decoder.h"
#include "raw_opus_decoder.h"
#include "audio_mem.h"
#include "audio_thread.h"
#include "board.h"
#include "audio_processor.h"

static char *TAG = "audio_processor";

#define audio_pipe_safe_free(x, fn) do { \
    if (x) {                             \
        fn(x);                           \
        x = NULL;                        \
    }                                    \
} while (0)

struct audio_recorder_s {
    audio_element_handle_t  i2s_reader;
    audio_element_handle_t  raw_stream;
    audio_element_handle_t  algo_stream;
    audio_pipeline_handle_t pipeline;
};

struct audio_player_s {
    audio_element_handle_t  i2s_writer;
    audio_element_handle_t  raw_stream;
    audio_element_handle_t  filter;
    audio_element_handle_t  opus_decoder_stream;
    audio_pipeline_handle_t pipeline;
    pipe_player_state_e     player_state;
};

typedef struct {
    audio_pipeline_handle_t pipeline;
    audio_element_handle_t  tone_stream_reader;
    audio_element_handle_t  i2s_stream_writer;
    audio_element_handle_t  mp3_decoder;
    audio_element_handle_t  filter;
    pipe_player_state_e     player_state;
    bool                    running;
    tone_play_callback_t    tone_cb;
} audio_player_t;

static          audio_player_t      *s_audio_player      = NULL;
static struct   audio_player_s      *s_player_pipeline   = NULL;

static int algo_read_data_callback(audio_element_handle_t self, char *buffer, int len, TickType_t ticks_to_wait, void *context)
{
    audio_element_handle_t i2s_reader = (audio_element_handle_t)context;
    return audio_element_input(i2s_reader, buffer, len);
}

audio_recorder_handle_t recorder_pipeline_open()
{
    struct audio_recorder_s *recorder = audio_calloc(1, sizeof(struct audio_recorder_s));
    if (recorder == NULL) {
        ESP_LOGE(TAG, "No mem for recorder");
        return NULL;
    }
    audio_pipeline_cfg_t pipeline_cfg = DEFAULT_AUDIO_PIPELINE_CONFIG();
    recorder->pipeline = audio_pipeline_init(&pipeline_cfg);
    AUDIO_MEM_CHECK(TAG, recorder->pipeline, goto _exit);

#if CONFIG_ESP32_S3_KORVO2_V3_BOARD || CONFIG_ESP32_S3_BOX_BOARD
    i2s_stream_cfg_t i2s_cfg = I2S_STREAM_CFG_DEFAULT_WITH_PARA(0, 16000, I2S_DATA_BIT_WIDTH_32BIT, AUDIO_STREAM_READER);
    i2s_stream_set_channel_type(&i2s_cfg, I2S_CHANNEL_TYPE_ONLY_LEFT);
#else
    i2s_stream_cfg_t i2s_cfg = I2S_STREAM_CFG_DEFAULT_WITH_PARA(0, 16000, I2S_DATA_BIT_WIDTH_16BIT, AUDIO_STREAM_READER);
#endif
    i2s_cfg.task_stack = -1;
    recorder->i2s_reader = i2s_stream_init(&i2s_cfg);
    AUDIO_MEM_CHECK(TAG, recorder->i2s_reader, goto _exit);

    algorithm_stream_cfg_t algo_config = ALGORITHM_STREAM_CFG_DEFAULT();
    algo_config.sample_rate = 16000;
    algo_config.out_rb_size = 26 * 1024;
    algo_config.task_core = 1;
#if CONFIG_ESP32_S3_KORVO2_V3_BOARD || CONFIG_ESP32_S3_BOX_BOARD
    algo_config.input_format = "RM";
#else
    algo_config.input_format = "MR";
#endif
    recorder->algo_stream = algo_stream_init(&algo_config);
    AUDIO_MEM_CHECK(TAG, recorder->algo_stream, goto _exit);
    audio_element_set_music_info(recorder->algo_stream, 16000, 1, 16);
    audio_element_set_read_cb(recorder->algo_stream, algo_read_data_callback, (void *)recorder->i2s_reader);
    audio_element_set_input_timeout(recorder->algo_stream, portMAX_DELAY);

    raw_stream_cfg_t raw_cfg = RAW_STREAM_CFG_DEFAULT();
    recorder->raw_stream = raw_stream_init(&raw_cfg);
    AUDIO_MEM_CHECK(TAG, recorder->raw_stream, goto _exit);

    audio_pipeline_register(recorder->pipeline, recorder->algo_stream, "algo_stream");
    audio_pipeline_register(recorder->pipeline, recorder->raw_stream, "raw_read");

    const char *link_tag[2] = {"algo_stream", "raw_read"};
    audio_pipeline_link(recorder->pipeline, &link_tag[0], sizeof(link_tag) / sizeof(char *));
    audio_pipeline_run(recorder->pipeline);
    return (audio_recorder_handle_t)recorder;
_exit:
    ESP_LOGE(TAG, "Failed to init recorder pipeline");
    if (recorder->i2s_reader) {
        audio_element_deinit(recorder->i2s_reader);
    }
    if (recorder->raw_stream) {
        audio_element_deinit(recorder->raw_stream);
    }
    if (recorder->algo_stream) {
        audio_element_deinit(recorder->algo_stream);
    }
    if (recorder->pipeline) {
        audio_pipeline_deinit(recorder->pipeline);
    }
    return NULL;
}

esp_err_t recorder_pipeline_run(audio_recorder_handle_t recorder)
{
    if (recorder == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    return audio_pipeline_run(recorder->pipeline);
}

esp_err_t recorder_pipeline_read(audio_recorder_handle_t recorder, char *buffer, int len)
{
    if (recorder == NULL || buffer == NULL || len <= 0) {
        return ESP_ERR_INVALID_ARG;
    }
    // ESP_LOGE(TAG, "recorder_pipeline_read");
    return raw_stream_read(recorder->raw_stream, buffer, len);
}

esp_err_t recorder_pipeline_stop(audio_recorder_handle_t recorder)
{
    if (recorder == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    audio_pipeline_stop(recorder->pipeline);
    audio_pipeline_wait_for_stop(recorder->pipeline);
    audio_pipeline_reset_elements(recorder->pipeline);
    audio_pipeline_reset_ringbuffer(recorder->pipeline);
    audio_pipeline_reset_items_state(recorder->pipeline);
    return ESP_OK;
}

esp_err_t recorder_pipeline_close(audio_recorder_handle_t recorder)
{
    if (recorder == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    audio_pipeline_terminate(recorder->pipeline);
    audio_pipeline_deinit(recorder->pipeline);
    return ESP_OK;
}

static int opus_audio_data_callback(audio_element_handle_t self, char *buffer, int len, TickType_t ticks_to_wait, void *context)
{
    audio_element_handle_t i2s_writer = (audio_element_handle_t)context;
    return audio_element_output(i2s_writer, buffer, len);
}

audio_player_handle_t player_pipeline_open()
{
    struct audio_player_s *player = audio_calloc(1, sizeof(struct audio_player_s));
    if (player == NULL) {
        ESP_LOGE(TAG, "No mem for player");
        return NULL;
    }
    player->player_state = PIPE_STATE_IDLE;
    s_player_pipeline = player;

    audio_pipeline_cfg_t pipeline_cfg = DEFAULT_AUDIO_PIPELINE_CONFIG();
    player->pipeline = audio_pipeline_init(&pipeline_cfg);
    AUDIO_MEM_CHECK(TAG, player->pipeline, goto _exit)
#if CONFIG_ESP32_S3_KORVO2_V3_BOARD || CONFIG_ESP32_S3_BOX_BOARD
    i2s_stream_cfg_t i2s_cfg = I2S_STREAM_CFG_DEFAULT_WITH_PARA(0, 16000, I2S_DATA_BIT_WIDTH_32BIT, AUDIO_STREAM_WRITER);
    i2s_stream_set_channel_type(&i2s_cfg, I2S_CHANNEL_TYPE_ONLY_LEFT);
    i2s_cfg.need_expand = true;
#else
    i2s_stream_cfg_t i2s_cfg = I2S_STREAM_CFG_DEFAULT_WITH_PARA(0, 16000, I2S_DATA_BIT_WIDTH_16BIT, AUDIO_STREAM_WRITER);
#endif
    i2s_cfg.task_stack = -1;
    player->i2s_writer = i2s_stream_init(&i2s_cfg);
    AUDIO_MEM_CHECK(TAG, player->i2s_writer, goto _exit);

    raw_stream_cfg_t raw_cfg = RAW_STREAM_CFG_DEFAULT();
    player->raw_stream = raw_stream_init(&raw_cfg);
    AUDIO_MEM_CHECK(TAG, player->raw_stream, goto _exit)

    raw_opus_dec_cfg_t opus_dec_cfg = RAW_OPUS_DEC_CONFIG_DEFAULT();
    opus_dec_cfg.enable_frame_length_prefix = true;
    opus_dec_cfg.sample_rate = 16000;
    opus_dec_cfg.channels = 1;
    opus_dec_cfg.task_core = 1;
    player->opus_decoder_stream = raw_opus_decoder_init(&opus_dec_cfg);
    AUDIO_MEM_CHECK(TAG, player->opus_decoder_stream, goto _exit);

    rsp_filter_cfg_t filter_cfg = DEFAULT_RESAMPLE_FILTER_CONFIG();
    filter_cfg.src_ch = 1;
    filter_cfg.src_rate = 16000;
#if CONFIG_ESP32_S3_KORVO2_V3_BOARD || CONFIG_ESP32_S3_BOX_BOARD
    filter_cfg.dest_ch = 1;
#else
    filter_cfg.dest_ch = 2;
#endif
    filter_cfg.dest_rate = 16000;
    filter_cfg.stack_in_ext = true;
    filter_cfg.task_core = 1;
    filter_cfg.complexity = 2;
    player->filter = rsp_filter_init(&filter_cfg);
    AUDIO_MEM_CHECK(TAG, player->filter, goto _exit);
    audio_element_set_write_cb(player->filter, opus_audio_data_callback, (void *)player->i2s_writer);

    audio_pipeline_register(player->pipeline, player->raw_stream, "raw_stream");
    audio_pipeline_register(player->pipeline, player->opus_decoder_stream, "raw_opus");
    audio_pipeline_register(player->pipeline, player->filter, "filter");

    const char *link_tag[3] = {"raw_stream", "raw_opus", "filter"};
    audio_pipeline_link(player->pipeline, &link_tag[0], 3);
    audio_pipeline_run(player->pipeline);
    return (audio_player_handle_t)player;
_exit:
    ESP_LOGE(TAG, "Failed to init player pipeline");
    if (player->i2s_writer) {
        audio_element_deinit(player->i2s_writer);
    }
    if (player->raw_stream) {
        audio_element_deinit(player->raw_stream);
    }
    if (player->opus_decoder_stream) {
        audio_element_deinit(player->opus_decoder_stream);
    }
    if (player->filter) {
        audio_element_deinit(player->filter);
    }
    if (player->pipeline) {
        audio_pipeline_deinit(player->pipeline);
    }
    return NULL;

}

static esp_err_t _player_i2s_write_cb(audio_element_handle_t self, char *buffer, int len, TickType_t ticks_to_wait, void *context)
{    
    return audio_element_output(s_player_pipeline->i2s_writer, buffer, len);
}

static esp_err_t _player_write_nop_cb(audio_element_handle_t self, char *buffer, int len, TickType_t ticks_to_wait, void *context)
{
    return len;
}

esp_err_t player_pipeline_run(void)
{
    ESP_RETURN_ON_FALSE(s_player_pipeline != NULL, ESP_FAIL, TAG, "player pipeline not initialized");
    if (s_player_pipeline->player_state == PIPE_STATE_RUNNING) {
        ESP_LOGW(TAG, "player pipe is already running state");
        return ESP_OK;
    }
    
    ESP_LOGI(TAG, "player pipe start running");
    audio_element_set_write_cb(s_player_pipeline->filter, _player_i2s_write_cb, NULL);
    s_player_pipeline->player_state = PIPE_STATE_RUNNING;
    return ESP_OK;
}

esp_err_t player_pipeline_stop(void)
{
    ESP_RETURN_ON_FALSE(s_player_pipeline != NULL, ESP_FAIL, TAG, "player pipeline not initialized");
    if (s_player_pipeline->player_state == PIPE_STATE_IDLE) {
        ESP_LOGW(TAG, "player pipe is idle state");
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "player pipe stop running");
    audio_element_set_write_cb(s_player_pipeline->filter, _player_write_nop_cb, NULL);
    s_player_pipeline->player_state = PIPE_STATE_IDLE;
    return ESP_OK;
}

esp_err_t player_pipeline_write(audio_player_handle_t player, char *buffer, int len)
{
    if (player == NULL || buffer == NULL || len <= 0) {
        return ESP_ERR_INVALID_ARG;
    }
    return raw_stream_write(player->raw_stream, buffer, len);
}

esp_err_t player_pipeline_close(audio_player_handle_t player)
{
    if (player == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    audio_pipeline_terminate(player->pipeline);
    audio_pipeline_deinit(player->pipeline);
    return ESP_OK;
}

esp_err_t audio_tone_play(tone_type_t tone_url)
{
    ESP_RETURN_ON_FALSE(s_audio_player != NULL, ESP_FAIL, TAG, "audio tone not initialized");
    if (s_audio_player->player_state == PIPE_STATE_RUNNING) {
        return ESP_FAIL;
    }
    ESP_LOGI(TAG, "audio_tone_play %d", tone_url);
    
    audio_element_set_uri(s_audio_player->tone_stream_reader, tone_uri[tone_url]);
    audio_pipeline_run(s_audio_player->pipeline);
    s_audio_player->player_state = PIPE_STATE_RUNNING;
    return ESP_OK;
}

esp_err_t audio_tone_stop(void)
{
    ESP_RETURN_ON_FALSE(s_audio_player != NULL, ESP_FAIL, TAG, "audio tone not initialized");
    if (s_audio_player->player_state == PIPE_STATE_IDLE) {
        return ESP_FAIL;
    }
    audio_pipeline_stop(s_audio_player->pipeline);
    audio_pipeline_wait_for_stop(s_audio_player->pipeline);
    audio_pipeline_terminate(s_audio_player->pipeline);
    audio_pipeline_reset_ringbuffer(s_audio_player->pipeline);
    audio_pipeline_reset_elements(s_audio_player->pipeline);
    s_audio_player->player_state = PIPE_STATE_IDLE;
    return ESP_OK;
}

static void audio_player_state_task(void *arg)
{
    audio_event_iface_handle_t evt = (audio_event_iface_handle_t) arg;

    s_audio_player->running = true;
    while (1) {
        audio_event_iface_msg_t msg;
        esp_err_t ret = audio_event_iface_listen(evt, &msg, portMAX_DELAY);
        if (ret != ESP_OK) {
            ESP_LOGE(TAG, "[ * ] Event interface error : %d", ret);
            continue;
        }
        if (msg.source_type == AUDIO_ELEMENT_TYPE_ELEMENT
            && msg.source == (void *) s_audio_player->mp3_decoder
            && msg.cmd == AEL_MSG_CMD_REPORT_MUSIC_INFO) {
            audio_element_info_t music_info = {0};
            audio_element_getinfo(s_audio_player->mp3_decoder, &music_info);
            ESP_LOGI(TAG, "[ * ] Receive music info from wav decoder, sample_rates=%d, bits=%d, ch=%d",
                    music_info.sample_rates, music_info.bits, music_info.channels);
            rsp_filter_set_src_info(s_audio_player->filter, music_info.sample_rates, music_info.channels);
            continue;
        }
        /* Stop when the last pipeline element receives stop event */
        if ((msg.source_type == AUDIO_ELEMENT_TYPE_ELEMENT && msg.source == (void *) s_audio_player->filter
            && msg.cmd == AEL_MSG_CMD_REPORT_STATUS)
            && ((int)msg.data == AEL_STATUS_STATE_FINISHED)) {
            ESP_LOGI(TAG, "[ * ] Stop event received");
            // To suppress unwanted transition noise
            vTaskDelay(pdMS_TO_TICKS(20));
            // Restore coze audio play
            player_pipeline_run();
            s_audio_player->tone_cb(AEL_STATUS_STATE_FINISHED);
        }
    }
}

static esp_err_t _player_write_mp3_cb(audio_element_handle_t self, char *buffer, int len, TickType_t ticks_to_wait, void *context)
{    
    return audio_element_output(s_player_pipeline->i2s_writer, buffer, len);
}

esp_err_t audio_tone_init(tone_play_callback_t callback)
{
    s_audio_player = (audio_player_t *)audio_calloc(1, sizeof(audio_player_t));
    AUDIO_MEM_CHECK(TAG, s_audio_player, goto _exit_open);

    ESP_LOGI(TAG, "Create audio pipeline for audio player");
    audio_pipeline_cfg_t pipeline_cfg = DEFAULT_AUDIO_PIPELINE_CONFIG();
    s_audio_player->pipeline = audio_pipeline_init(&pipeline_cfg);
    AUDIO_MEM_CHECK(TAG, s_audio_player->pipeline, goto _exit_open);

    ESP_LOGI(TAG, "Create tone stream to read data from flash");
    tone_stream_cfg_t tone_cfg = TONE_STREAM_CFG_DEFAULT();
    tone_cfg.type = AUDIO_STREAM_READER;
    s_audio_player->tone_stream_reader = tone_stream_init(&tone_cfg);
    AUDIO_NULL_CHECK(TAG, s_audio_player->tone_stream_reader, goto _exit_open);

    ESP_LOGI(TAG, "Create mp3 decoder to decode mp3 file");
    mp3_decoder_cfg_t mp3_cfg = DEFAULT_MP3_DECODER_CONFIG();
    s_audio_player->mp3_decoder = mp3_decoder_init(&mp3_cfg);
    AUDIO_NULL_CHECK(TAG, s_audio_player->mp3_decoder, goto _exit_open);

    ESP_LOGI(TAG, "Create filter for mp3");
    rsp_filter_cfg_t filter_cfg = DEFAULT_RESAMPLE_FILTER_CONFIG();
    filter_cfg.src_ch = 1;
    filter_cfg.src_rate = 16000;
#if CONFIG_ESP32_S3_KORVO2_V3_BOARD || CONFIG_ESP32_S3_BOX_BOARD
    filter_cfg.dest_ch = 1;
#else
    filter_cfg.dest_ch = 2;
#endif
    filter_cfg.dest_rate = 16000;
    filter_cfg.stack_in_ext = true;
    filter_cfg.task_core = 1;
    filter_cfg.complexity = 2;
    s_audio_player->filter = rsp_filter_init(&filter_cfg);
    AUDIO_MEM_CHECK(TAG, s_audio_player->filter, goto _exit_open);
    audio_element_set_write_cb(s_audio_player->filter, _player_write_mp3_cb, NULL);

    ESP_LOGI(TAG, "Register all elements to audio pipeline");
    audio_pipeline_register(s_audio_player->pipeline, s_audio_player->tone_stream_reader, "tone");
    audio_pipeline_register(s_audio_player->pipeline, s_audio_player->mp3_decoder, "mp3");
    audio_pipeline_register(s_audio_player->pipeline, s_audio_player->filter, "filter");

    ESP_LOGI(TAG, "Link it together [flash]-->tone_stream-->mp3_decoder-->filter-->[codec_chip]");
    const char *link_tag[3] = {"tone", "mp3", "filter"};
    audio_pipeline_link(s_audio_player->pipeline, &link_tag[0], 3);


    esp_periph_config_t periph_cfg = DEFAULT_ESP_PERIPH_SET_CONFIG();
    esp_periph_set_handle_t set = esp_periph_set_init(&periph_cfg);
    audio_event_iface_cfg_t evt_cfg = AUDIO_EVENT_IFACE_DEFAULT_CFG();
    audio_event_iface_handle_t evt = audio_event_iface_init(&evt_cfg);
    audio_pipeline_set_listener(s_audio_player->pipeline, evt);
    audio_event_iface_set_listener(esp_periph_set_get_event_iface(set), evt);

    s_audio_player->tone_cb = callback;
    audio_thread_create(NULL, "audio_player_state_task", audio_player_state_task, (void *)evt, 5 * 1024, 15, true, 1);

    return ESP_OK;

_exit_open:
    audio_pipe_safe_free(s_audio_player->tone_stream_reader, audio_element_deinit);
    audio_pipe_safe_free(s_audio_player->mp3_decoder, audio_element_deinit);
    audio_pipe_safe_free(s_audio_player->filter, audio_element_deinit);
    audio_pipe_safe_free(s_audio_player->pipeline, audio_pipeline_deinit);
    audio_pipe_safe_free(s_audio_player, audio_free);
    return ESP_FAIL;
}