#include "system.h"

#include "audio_codec.h"
#include "board.h"

#include <esp_log.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>

#include <esp_audio_simple_dec.h>
#include <esp_audio_simple_dec_default.h>
#include <esp_audio_dec_default.h>
#include <esp_ae_rate_cvt.h>

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>

#define TAG "sdcard_mp3"
#define MP3_READ_BUFFER_SIZE 512
#define MP3_OUTPUT_BUFFER_SIZE 4096
#define MP3_PLAYER_STACK_SIZE 8192
#define MP3_PLAYER_PRIORITY 5

typedef struct {
    char path[SDCARD_PATH_MAX_LEN];
    float speed;
} mp3_player_task_arg_t;

static TaskHandle_t mp3_player_task_handle;
static volatile bool mp3_stop_requested;
static volatile bool mp3_pause_requested;

static std::vector<int16_t> convert_channels(const int16_t *input, size_t frame_count, uint8_t src_channels, uint8_t dst_channels)
{
    std::vector<int16_t> output(frame_count * dst_channels);
    if (src_channels == dst_channels) {
        std::memcpy(output.data(), input, output.size() * sizeof(int16_t));
        return output;
    }

    for (size_t frame = 0; frame < frame_count; ++frame) {
        if (dst_channels == 1) {
            int32_t mixed = 0;
            for (uint8_t ch = 0; ch < src_channels; ++ch) {
                mixed += input[frame * src_channels + ch];
            }
            output[frame] = static_cast<int16_t>(mixed / src_channels);
        } else if (src_channels == 1) {
            for (uint8_t ch = 0; ch < dst_channels; ++ch) {
                output[frame * dst_channels + ch] = input[frame];
            }
        } else {
            for (uint8_t ch = 0; ch < dst_channels; ++ch) {
                output[frame * dst_channels + ch] = input[frame * src_channels + std::min<uint8_t>(ch, src_channels - 1)];
            }
        }
    }

    return output;
}

static bool is_supported_speed(float speed)
{
    const float supported_speeds[] = {0.5f, 0.75f, 1.0f, 1.25f, 1.5f, 2.0f};
    for (size_t i = 0; i < sizeof(supported_speeds) / sizeof(supported_speeds[0]); ++i) {
        float diff = speed - supported_speeds[i];
        if (diff < 0) {
            diff = -diff;
        }
        if (diff < 0.001f) {
            return true;
        }
    }
    return false;
}

static bool wait_if_paused()
{
    while (mp3_pause_requested && !mp3_stop_requested) {
        vTaskDelay(pdMS_TO_TICKS(20));
    }
    return !mp3_stop_requested;
}

static bool output_pcm(AudioCodec *codec, esp_ae_rate_cvt_handle_t rate_cvt_handle, const uint8_t *buffer,
    uint32_t size, uint8_t src_channels, uint32_t src_rate)
{
    uint8_t dst_channels = codec->output_channels();
    uint32_t dst_rate = codec->output_sample_rate();
    size_t frame_count = size / (sizeof(int16_t) * src_channels);
    if (frame_count == 0) {
        return true;
    }

    auto channel_converted = convert_channels(reinterpret_cast<const int16_t *>(buffer), frame_count, src_channels, dst_channels);
    if (src_rate == dst_rate || !rate_cvt_handle) {
        codec->OutputData(channel_converted);
        return true;
    }

    uint32_t max_out_frames = 0;
    if (esp_ae_rate_cvt_get_max_out_sample_num(rate_cvt_handle, frame_count, &max_out_frames) != ESP_AE_ERR_OK) {
        ESP_LOGE(TAG, "get resample output size failed");
        return false;
    }

    std::vector<int16_t> resampled(max_out_frames * dst_channels);
    uint32_t actual_out_frames = max_out_frames;
    if (esp_ae_rate_cvt_process(rate_cvt_handle, reinterpret_cast<esp_ae_sample_t>(channel_converted.data()), frame_count,
            reinterpret_cast<esp_ae_sample_t>(resampled.data()), &actual_out_frames) != ESP_AE_ERR_OK) {
        ESP_LOGE(TAG, "resample MP3 PCM failed");
        return false;
    }

    resampled.resize(actual_out_frames * dst_channels);
    codec->OutputData(resampled);
    return true;
}

static void mp3_player_task(void *arg)
{
    mp3_player_task_handle = xTaskGetCurrentTaskHandle();
    mp3_player_task_arg_t *task_arg = static_cast<mp3_player_task_arg_t *>(arg);
    char path[SDCARD_PATH_MAX_LEN];
    strlcpy(path, task_arg->path, sizeof(path));
    float speed = task_arg->speed;
    free(task_arg);

    FILE *fp = fopen(path, "rb");
    if (!fp) {
        ESP_LOGE(TAG, "open MP3 failed: %s", path);
        mp3_player_task_handle = nullptr;
        mp3_stop_requested = false;
        vTaskDelete(nullptr);
        return;
    }

    auto codec = Board::GetInstance().GetAudioCodec();
    if (!codec) {
        ESP_LOGE(TAG, "audio codec unavailable");
        fclose(fp);
        mp3_player_task_handle = nullptr;
        mp3_stop_requested = false;
        vTaskDelete(nullptr);
        return;
    }

    esp_audio_dec_register_default();
    esp_audio_simple_dec_register_default();

    esp_audio_simple_dec_handle_t decoder = nullptr;
    esp_audio_simple_dec_cfg_t dec_cfg = {
        .dec_type = ESP_AUDIO_SIMPLE_DEC_TYPE_MP3,
        .dec_cfg = nullptr,
        .cfg_size = 0,
        .use_frame_dec = false,
    };

    esp_audio_err_t ret = esp_audio_simple_dec_open(&dec_cfg, &decoder);
    if (ret != ESP_AUDIO_ERR_OK) {
        ESP_LOGE(TAG, "open MP3 decoder failed: %d", ret);
        esp_audio_simple_dec_unregister_default();
        esp_audio_dec_unregister_default();
        fclose(fp);
        mp3_player_task_handle = nullptr;
        mp3_stop_requested = false;
        vTaskDelete(nullptr);
        return;
    }

    uint8_t *in_buf = static_cast<uint8_t *>(malloc(MP3_READ_BUFFER_SIZE));
    uint8_t *out_buf = static_cast<uint8_t *>(malloc(MP3_OUTPUT_BUFFER_SIZE));
    uint32_t out_buf_size = MP3_OUTPUT_BUFFER_SIZE;
    bool output_enabled_before = codec->output_enabled();
    bool audio_info_checked = false;
    uint32_t mp3_sample_rate = 0;
    uint8_t mp3_channels = 0;
    esp_ae_rate_cvt_handle_t rate_cvt_handle = nullptr;

    if (!in_buf || !out_buf) {
        ESP_LOGE(TAG, "no memory for MP3 playback");
        ret = ESP_AUDIO_ERR_MEM_LACK;
    } else {
        ESP_LOGI(TAG, "play MP3: %s, speed: %.2fx", path, speed);

        if (wait_if_paused()) {
            codec->EnableOutput(true);
        }

        while (ret == ESP_AUDIO_ERR_OK && !mp3_stop_requested) {
            if (!wait_if_paused()) {
                break;
            }

            int read_size = fread(in_buf, 1, MP3_READ_BUFFER_SIZE, fp);
            if (read_size <= 0) {
                break;
            }

            esp_audio_simple_dec_raw_t raw = {
                .buffer = in_buf,
                .len = static_cast<uint32_t>(read_size),
                .eos = read_size < MP3_READ_BUFFER_SIZE,
                .consumed = 0,
                .frame_recover = ESP_AUDIO_SIMPLE_DEC_RECOVERY_NONE,
            };

            while (raw.len > 0 && !mp3_stop_requested) {
                if (!wait_if_paused()) {
                    break;
                }

                esp_audio_simple_dec_out_t out_frame = {
                    .buffer = out_buf,
                    .len = out_buf_size,
                    .needed_size = 0,
                    .decoded_size = 0,
                };

                ret = esp_audio_simple_dec_process(decoder, &raw, &out_frame);
                if (ret == ESP_AUDIO_ERR_BUFF_NOT_ENOUGH) {
                    uint8_t *new_buf = static_cast<uint8_t *>(realloc(out_buf, out_frame.needed_size));
                    if (!new_buf) {
                        ESP_LOGE(TAG, "resize MP3 output buffer failed");
                        ret = ESP_AUDIO_ERR_MEM_LACK;
                        break;
                    }
                    out_buf = new_buf;
                    out_buf_size = out_frame.needed_size;
                    continue;
                }

                if (ret != ESP_AUDIO_ERR_OK) {
                    ESP_LOGE(TAG, "decode MP3 failed: %d", ret);
                    break;
                }

                if (out_frame.decoded_size > 0) {
                    if (!audio_info_checked) {
                        esp_audio_simple_dec_info_t info = {};
                        if (esp_audio_simple_dec_get_info(decoder, &info) == ESP_AUDIO_ERR_OK) {
                            ESP_LOGI(TAG, "MP3 info: %lu Hz, %u bits, %u ch", info.sample_rate, info.bits_per_sample, info.channel);
                            if (info.bits_per_sample != 16 || info.channel == 0 || codec->output_channels() == 0) {
                                ESP_LOGE(TAG, "MP3 format unsupported: %lu Hz, %u bits, %u ch", info.sample_rate, info.bits_per_sample, info.channel);
                                ret = ESP_AUDIO_ERR_NOT_SUPPORT;
                                break;
                            }
                            mp3_sample_rate = static_cast<uint32_t>(info.sample_rate * speed);
                            mp3_channels = info.channel;
                            if (mp3_sample_rate != static_cast<uint32_t>(codec->output_sample_rate())) {
                                esp_ae_rate_cvt_cfg_t rate_cvt_cfg = {
                                    .src_rate = mp3_sample_rate,
                                    .dest_rate = static_cast<uint32_t>(codec->output_sample_rate()),
                                    .channel = static_cast<uint8_t>(codec->output_channels()),
                                    .bits_per_sample = 16,
                                    .complexity = 2,
                                    .perf_type = ESP_AE_RATE_CVT_PERF_TYPE_SPEED,
                                };
                                if (esp_ae_rate_cvt_open(&rate_cvt_cfg, &rate_cvt_handle) != ESP_AE_ERR_OK || !rate_cvt_handle) {
                                    ESP_LOGE(TAG, "open MP3 resampler failed");
                                    ret = ESP_AUDIO_ERR_NOT_SUPPORT;
                                    break;
                                }
                            }
                        }
                        audio_info_checked = true;
                    }

                    if (!output_pcm(codec, rate_cvt_handle, out_frame.buffer, out_frame.decoded_size, mp3_channels, mp3_sample_rate)) {
                        ret = ESP_AUDIO_ERR_FAIL;
                        break;
                    }
                }

                if (raw.consumed == 0) {
                    break;
                }
                raw.len -= raw.consumed;
                raw.buffer += raw.consumed;
                raw.consumed = 0;
            }

            if (raw.eos || ret != ESP_AUDIO_ERR_OK) {
                break;
            }
        }
    }

    if (rate_cvt_handle) {
        esp_ae_rate_cvt_close(rate_cvt_handle);
    }
    if (!output_enabled_before) {
        codec->EnableOutput(false);
    }
    free(in_buf);
    free(out_buf);
    esp_audio_simple_dec_close(decoder);
    esp_audio_simple_dec_unregister_default();
    esp_audio_dec_unregister_default();
    fclose(fp);

    ESP_LOGI(TAG, "MP3 playback %s: %s", mp3_stop_requested ? "stopped" : "finished", path);
    mp3_player_task_handle = nullptr;
    mp3_stop_requested = false;
    mp3_pause_requested = false;
    vTaskDelete(nullptr);
}

extern "C" esp_err_t SDCard_Play_MP3_File(const char *path)
{
    return SDCard_Play_MP3_File_With_Speed(path, 1.0f);
}

extern "C" esp_err_t SDCard_Play_MP3_File_With_Speed(const char *path, float speed)
{
    if (!path || path[0] == '\0' || !is_supported_speed(speed)) {
        return ESP_ERR_INVALID_ARG;
    }

    bool start_paused = mp3_pause_requested;
    if (mp3_player_task_handle) {
        mp3_stop_requested = true;
        mp3_pause_requested = false;
        while (mp3_player_task_handle) {
            vTaskDelay(pdMS_TO_TICKS(20));
        }
    }

    mp3_player_task_arg_t *arg = static_cast<mp3_player_task_arg_t *>(calloc(1, sizeof(mp3_player_task_arg_t)));
    if (!arg) {
        return ESP_ERR_NO_MEM;
    }
    strlcpy(arg->path, path, sizeof(arg->path));
    arg->speed = speed;
    mp3_stop_requested = false;
    mp3_pause_requested = start_paused;

    BaseType_t ret = xTaskCreate(mp3_player_task, "mp3_player", MP3_PLAYER_STACK_SIZE, arg, MP3_PLAYER_PRIORITY, &mp3_player_task_handle);
    if (ret != pdPASS) {
        free(arg);
        mp3_player_task_handle = nullptr;
        return ESP_FAIL;
    }

    return ESP_OK;
}

extern "C" void SDCard_Stop_Play(void)
{
    mp3_pause_requested = false;
    if (!mp3_player_task_handle) {
        ESP_LOGW(TAG, "MP3 player is not running");
        return;
    }

    mp3_stop_requested = true;
}

extern "C" void SDCard_Pause_Play(void)
{
    mp3_pause_requested = true;
}

extern "C" void SDCard_Resume_Play(void)
{
    mp3_pause_requested = false;
}

extern "C" bool SDCard_Is_Playing(void)
{
    return mp3_player_task_handle != nullptr;
}

extern "C" bool SDCard_Is_Paused(void)
{
    return mp3_player_task_handle != nullptr && mp3_pause_requested;
}