/**
 ****************************************************************************************************
 * @file        audio_source.c
 * @author      HiFi Player Project
 * @version     V1.5
 * @date        2026-01-01
 * @brief       统一音源实现（WAV 解析 + FLAC 解码转发）
 *
 * 两种格式都输出「立体声 + 源文件原生位深」的 PCM，不做降位：
 *
 *   WAV  ：f_read 原始 PCM ──▶ 单声道则复制 ──▶ 输出
 *   FLAC ：flac_dec_read()  ──▶ 已在内部做声道规整 + 原生位深输出
 *
 * 格式识别：读取文件头 4 字节
 *   "RIFF" → WAV
 *   "fLaC" → 原生 FLAC
 *   "OggS" → Ogg FLAC
 ****************************************************************************************************
 */

#include "audio_source.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "esp_log.h"

#include "ff.h"
#include "flac_dec.h"

#define TAG                 "audio_src"

#define WAV_MAX_BLOCKALIGN  8       /* 支持的 WAV 帧对齐上限（16/24/32bit × 最多 2 声道） */
#define WAV_MAX_FRAMES      1024    /* WAV 原始 PCM 中转缓冲容纳的最大帧数 */

#define AUDIO_SRC_MAX_FRAMES 1024   /* 单次 audio_source_read 最多输出的帧数（与播放器缓冲一致） */
#define DECIM_CHUNK_FRAMES   256    /* 抽取时每次处理多少输出帧（控制临时缓冲大小） */
#define AUDIO_MAX_OUT_RATE   48000  /* 输出采样率上限 */
#define AUDIO_MAX_DECIM      8      /* 最大抽取倍数 */

/* ---------------- 内部结构 ---------------- */
struct audio_source {
    audio_codec_t codec;
    audio_info_t  info;

    uint64_t frames_read;       /* 已输出帧数（用于计算进度） */

    /* ---- WAV ---- */
    FIL      wav_fp;
    bool     wav_open;
    uint16_t block_align;       /* 源每帧字节数 = channels * out_bytes */
    uint32_t byte_rate;
    uint32_t data_offset;
    uint32_t data_size;
    uint32_t data_read;         /* 已读字节数 */
    uint8_t *scratch;           /* 原始 PCM 中转缓冲 */
    int      scratch_frames;    /* 缓冲能容纳的帧数 */

    /* ---- FLAC ---- */
    flac_dec_t *flac;

    /* ---- 高采样率抽取（192k/96k → ≤48k） ---- */
    uint8_t  decim;             /* 抽取倍数，1 表示不抽取 */
    uint8_t *decim_buf;         /* 抽取用的原生采样中转缓冲 */
};

/* ---------------- 位宽映射 ---------------- */

/**
 * @brief  源位深 → I2S 输出位宽
 */
static inline uint16_t pick_out_bits(uint16_t bits)
{
    if (bits <= 16)
    {
        return 16;
    }

    if (bits <= 24)
    {
        return 24;
    }

    return 32;
}

/* 采样率限制 / 抽取（实现见文件末尾，供 audio_source_open 使用） */
static bool decim_setup(audio_source_t *src);
static void decim_mix(const uint8_t *src, uint8_t *dst, int frames, int decim, uint8_t ob);

/* ---------------- WAV 解析 ---------------- */

static inline uint16_t rd_u16(const uint8_t *p)
{
    return (uint16_t)(p[0] | (p[1] << 8));
}

static inline uint32_t rd_u32(const uint8_t *p)
{
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

/**
 * @brief  解析 WAV 头部（支持 fmt/data 之间夹着 LIST、fact 等块）
 * @note   调用前文件指针需在文件开头；结束后文件指针位于 data 起始处
 */
static bool wav_parse(FIL *fp, audio_info_t *info, uint32_t *data_offset,
                      uint32_t *data_size, uint16_t *block_align, uint32_t *byte_rate,
                      uint16_t *bits, uint16_t *channels)
{
    uint8_t header[12];
    uint8_t chunk[8];
    uint8_t fmt[40];
    UINT br = 0;
    bool have_fmt = false;
    bool have_data = false;
    uint16_t audio_format = 0;

    if ((f_read(fp, header, sizeof(header), &br) != FR_OK) || (br != sizeof(header)))
    {
        return false;
    }

    if ((memcmp(header, "RIFF", 4) != 0) || (memcmp(header + 8, "WAVE", 4) != 0))
    {
        return false;
    }

    while (!have_data)
    {
        if ((f_read(fp, chunk, sizeof(chunk), &br) != FR_OK) || (br != sizeof(chunk)))
        {
            break;
        }

        uint32_t size = rd_u32(chunk + 4);

        if (memcmp(chunk, "fmt ", 4) == 0)
        {
            if ((size < 16) || (size > sizeof(fmt)))
            {
                break;
            }

            if ((f_read(fp, fmt, size, &br) != FR_OK) || (br != size))
            {
                break;
            }

            audio_format      = rd_u16(fmt);
            *channels         = rd_u16(fmt + 2);
            info->sample_rate = rd_u32(fmt + 4);
            *byte_rate        = rd_u32(fmt + 8);
            *block_align      = rd_u16(fmt + 12);
            *bits             = rd_u16(fmt + 14);

            if ((audio_format == 0xFFFE) && (size >= 26))
            {
                /* WAVE_FORMAT_EXTENSIBLE：真实格式在 SubFormat GUID 的前两个字节 */
                audio_format = rd_u16(fmt + 24);
            }

            have_fmt = true;
        }
        else if (memcmp(chunk, "data", 4) == 0)
        {
            FSIZE_t offset = f_tell(fp);
            FSIZE_t remain = f_size(fp) - offset;

            *data_offset = (uint32_t)offset;
            /* 有些录音文件 data 长度写 0 或 0xFFFFFFFF，此时用文件剩余长度 */
            *data_size = ((size == 0) || (size > (uint32_t)remain)) ? (uint32_t)remain : size;

            have_data = true;
            break;
        }
        else
        {
            /* 跳过不关心的块（块长度 2 字节对齐） */
            if (f_lseek(fp, f_tell(fp) + size + (size & 0x01)) != FR_OK)
            {
                break;
            }
        }
    }

    if (!have_fmt || !have_data)
    {
        return false;
    }

    if (*block_align == 0)
    {
        *block_align = (uint16_t)(((*bits) / 8) * (*channels));
    }

    if (*byte_rate == 0)
    {
        *byte_rate = info->sample_rate * (*block_align);
    }

    /* 只接受线性 PCM，位深需为 16/24/32（非整字节位深无法从 WAV 头可靠判断） */
    return (audio_format == 1) &&
           (*channels >= 1) && (*channels <= 2) &&
           (*block_align > 0) && (*block_align <= WAV_MAX_BLOCKALIGN) &&
           ((*bits == 16) || (*bits == 24) || (*bits == 32)) &&
           (info->sample_rate >= 8000) && (*byte_rate > 0);
}

/**
 * @brief  WAV：把原始 PCM 规整成立体声（保留原生位深）
 */
static int wav_read(audio_source_t *src, void *out, int max_frames)
{
    if ((src->data_read >= src->data_size) || (max_frames <= 0))
    {
        return 0;
    }

    uint32_t remain_frames = (src->data_size - src->data_read) / src->block_align;
    uint32_t frames = (remain_frames < (uint32_t)max_frames) ? remain_frames : (uint32_t)max_frames;

    if ((int)frames > src->scratch_frames)
    {
        frames = (uint32_t)src->scratch_frames;
    }

    if (frames == 0)
    {
        return 0;
    }

    UINT br = 0;
    size_t want = (size_t)frames * src->block_align;

    if (f_read(&src->wav_fp, src->scratch, (UINT)want, &br) != FR_OK)
    {
        return -1;
    }

    if (br < src->block_align)
    {
        return 0;                       /* 不足一帧，视作结束 */
    }

    frames = br / src->block_align;
    src->data_read += br;

    /* ---- 声道规整（位深保持不变） ---- */
    const uint8_t *s = src->scratch;
    uint8_t *dst = (uint8_t *)out;
    const size_t ob = src->info.out_bytes;

    if (src->info.channels == 2)
    {
        /* 最常见路径：立体声直接拷贝（block_align == 2 * out_bytes） */
        memcpy(dst, s, (size_t)frames * 2 * ob);
    }
    else
    {
        /* 单声道：左右都写同一样本 */
        for (uint32_t i = 0; i < frames; i++)
        {
            uint8_t *o = dst + i * 2 * ob;

            memcpy(o, s + i * ob, ob);
            memcpy(o + ob, s + i * ob, ob);
        }
    }

    /* 注意：frames_read 统一由 audio_source_read() 累加，这里不要重复计数 */
    return (int)frames;
}

/* ---------------- 对外接口 ---------------- */

audio_source_t *audio_source_open(const char *path, audio_info_t *info)
{
    if (path == NULL)
    {
        return NULL;
    }

    audio_source_t *src = (audio_source_t *)calloc(1, sizeof(audio_source_t));

    if (src == NULL)
    {
        return NULL;
    }

    /* ---------------- 1. 打开文件并嗅探文件头 ---------------- */
    if (f_open(&src->wav_fp, path, FA_READ) != FR_OK)
    {
        ESP_LOGE(TAG, "open %s failed", path);
        audio_source_close(src);
        return NULL;
    }
    src->wav_open = true;

    uint8_t magic[4] = {0};
    UINT br = 0;
    bool is_flac = false;

    if ((f_read(&src->wav_fp, magic, sizeof(magic), &br) == FR_OK) && (br == sizeof(magic)))
    {
        is_flac = (memcmp(magic, "fLaC", 4) == 0) || (memcmp(magic, "OggS", 4) == 0);
    }

    /* ---------------- 2A. FLAC ---------------- */
    if (is_flac)
    {
        /* 交给专门的解码器重新打开（它需要从文件开头开始解析） */
        f_close(&src->wav_fp);
        src->wav_open = false;

        flac_info_t fi;

        src->flac = flac_dec_open(path);

        if ((src->flac == NULL) || (flac_dec_get_info(src->flac, &fi) != 0))
        {
            ESP_LOGE(TAG, "flac open failed: %s", path);
            audio_source_close(src);
            return NULL;
        }

        src->codec = AUDIO_CODEC_FLAC;
        src->info.codec       = AUDIO_CODEC_FLAC;
        src->info.sample_rate = fi.sample_rate;
        src->info.bits        = fi.bits;
        src->info.channels    = fi.channels;
        src->info.out_bytes   = fi.out_bytes;
        src->info.out_bits    = (uint16_t)fi.out_bytes * 8;
        src->info.total_sec   = (fi.total_samples > 0) ?
                                (uint32_t)(fi.total_samples / fi.sample_rate) : 0;

        /* 高采样率（96k/192k）抽取到 ES8388 / MCLK 能承受的范围 */
        if (!decim_setup(src))
        {
            ESP_LOGE(TAG, "unsupported sample rate %u Hz: %s",
                     (unsigned)fi.sample_rate, path);
            audio_source_close(src);
            return NULL;
        }

        if (info != NULL)
        {
            *info = src->info;
        }

        ESP_LOGI(TAG, "%s: FLAC %u Hz%s, %u bit -> I2S %u bit @ %u Hz, %u ch, %u s",
                 path, (unsigned)fi.sample_rate,
                 (src->decim > 1) ? " (downsampled)" : "",
                 src->info.bits, src->info.out_bits,
                 (unsigned)src->info.sample_rate, src->info.channels,
                 (unsigned)src->info.total_sec);

        return src;
    }

    /* ---------------- 2B. WAV ---------------- */
    src->codec = AUDIO_CODEC_WAV;
    src->info.codec = AUDIO_CODEC_WAV;

    f_lseek(&src->wav_fp, 0);           /* 回到文件开头重新解析 */

    uint16_t bits = 0;
    uint16_t channels = 0;

    if (!wav_parse(&src->wav_fp, &src->info, &src->data_offset, &src->data_size,
                   &src->block_align, &src->byte_rate, &bits, &channels))
    {
        ESP_LOGW(TAG, "unsupported wav: %s", path);
        audio_source_close(src);
        return NULL;
    }

    src->info.bits      = bits;
    src->info.channels  = channels;
    src->info.out_bits  = pick_out_bits(bits);
    src->info.out_bytes = (uint8_t)(bits / 8);      /* WAV 位深必为 16/24/32 */
    src->info.total_sec = src->data_size / src->byte_rate;

    /* 高采样率（96k/192k）抽取到 ES8388 / MCLK 能承受的范围 */
    if (!decim_setup(src))
    {
        ESP_LOGW(TAG, "unsupported sample rate %u Hz: %s",
                 (unsigned)src->info.sample_rate, path);
        audio_source_close(src);
        return NULL;
    }

    /* 原始 PCM 中转缓冲（16/24/32bit × 最多 2 声道 → 每帧最多 8 字节） */
    src->scratch = (uint8_t *)malloc((size_t)WAV_MAX_FRAMES * src->block_align);

    if (src->scratch == NULL)
    {
        ESP_LOGE(TAG, "no memory for wav scratch");
        audio_source_close(src);
        return NULL;
    }
    src->scratch_frames = WAV_MAX_FRAMES;

    f_lseek(&src->wav_fp, src->data_offset);

    if (info != NULL)
    {
        *info = src->info;
    }

    ESP_LOGI(TAG, "%s: WAV %u Hz, %u bit -> I2S %u bit, %u ch, %u s",
             path, (unsigned)src->info.sample_rate, src->info.bits,
             src->info.out_bits, src->info.channels,
             (unsigned)src->info.total_sec);

    return src;
}

/**
 * @brief  带抽取的读取：把 decim 个原生帧合成 1 个输出帧
 */
static int audio_source_read_decim(audio_source_t *src, void *out, int max_frames)
{
    uint8_t *dst = (uint8_t *)out;
    const uint8_t ob = src->info.out_bytes;
    const int decim = src->decim;
    int produced = 0;

    while (produced < max_frames)
    {
        int room = max_frames - produced;

        if (room > DECIM_CHUNK_FRAMES)
        {
            room = DECIM_CHUNK_FRAMES;
        }

        if ((room * decim) > WAV_MAX_FRAMES)
        {
            room = WAV_MAX_FRAMES / decim;      /* 受中转缓冲帧数限制 */
        }

        if (room <= 0)
        {
            break;
        }

        /* 先凑满 room * decim 个原生帧，避免每次少读几帧造成丢样本 */
        const int need = room * decim;
        int have = 0;
        bool eof = false;

        while (have < need)
        {
            int got = (src->codec == AUDIO_CODEC_FLAC)
                          ? flac_dec_read(src->flac,
                                          src->decim_buf + (size_t)have * 2u * ob,
                                          need - have)
                          : wav_read(src,
                                     src->decim_buf + (size_t)have * 2u * ob,
                                     need - have);

            if (got < 0)
            {
                return -1;
            }

            if (got == 0)
            {
                eof = true;
                break;
            }

            have += got;
        }

        int frames = have / decim;

        if (frames > 0)
        {
            decim_mix(src->decim_buf, dst + (size_t)produced * 2u * ob,
                      frames, decim, ob);
            produced += frames;
        }

        if (eof)
        {
            break;                              /* 文件末尾 */
        }
    }

    if (produced > 0)
    {
        src->frames_read += (uint64_t)produced;
    }

    return produced;
}

int audio_source_read(audio_source_t *src, void *out, int max_frames)
{
    if ((src == NULL) || (out == NULL) || (max_frames <= 0))
    {
        return -1;
    }

    if (max_frames > AUDIO_SRC_MAX_FRAMES)
    {
        max_frames = AUDIO_SRC_MAX_FRAMES;      /* 保护调用者缓冲 */
    }

    if (src->decim > 1)
    {
        return audio_source_read_decim(src, out, max_frames);
    }

    int n = (src->codec == AUDIO_CODEC_FLAC)
                ? flac_dec_read(src->flac, out, max_frames)
                : wav_read(src, out, max_frames);

    if (n > 0)
    {
        src->frames_read += (uint64_t)n;
    }

    return n;
}

uint32_t audio_source_elapsed_sec(const audio_source_t *src)
{
    if ((src == NULL) || (src->info.sample_rate == 0))
    {
        return 0;
    }

    uint32_t sec = (uint32_t)(src->frames_read / src->info.sample_rate);

    if ((src->info.total_sec > 0) && (sec > src->info.total_sec))
    {
        sec = src->info.total_sec;
    }

    return sec;
}

void audio_source_close(audio_source_t *src)
{
    if (src == NULL)
    {
        return;
    }

    if (src->flac != NULL)
    {
        flac_dec_close(src->flac);
        src->flac = NULL;
    }

    if (src->wav_open)
    {
        f_close(&src->wav_fp);
        src->wav_open = false;
    }

    free(src->scratch);
    src->scratch = NULL;

    free(src->decim_buf);
    src->decim_buf = NULL;

    free(src);
}

/* ---------------- 采样率限制 ---------------- */

/**
 * @brief  计算抽取倍数（把采样率压到 AUDIO_MAX_OUT_RATE 以内）
 * @note   板载 ES8388 的 MCLK 由 I2S 产生，MCLK = mclk_multiple × fs.
 *         24bit 时 mclk_multiple 必须是 384（3 的倍数），若直接播放 192kHz,
 *         MCLK 会高达 73.7MHz，远超 ES8388 的输入范围，结果是无声音或杂音.
 *         因此这里对 88.2k/96k/192k 等做整数抽取，统一落到 ≤48kHz.
 */
static uint8_t pick_decim(uint32_t rate)
{
    uint8_t d = 1;

    while (((rate / d) > AUDIO_MAX_OUT_RATE) && (d < AUDIO_MAX_DECIM))
    {
        d++;
    }

    return d;
}

/**
 * @brief  按输出采样率为音源分配抽取缓冲
 * @retval true 成功；false 采样率过高或内存不足
 */
static bool decim_setup(audio_source_t *src)
{
    uint8_t d = pick_decim(src->info.sample_rate);

    if ((src->info.sample_rate / d) > AUDIO_MAX_OUT_RATE)
    {
        return false;                       /* 采样率过高，无法用整数抽取覆盖 */
    }

    src->decim = d;

    if (d <= 1)
    {
        return true;                        /* 常见 44.1k/48k：完全不额外占用内存 */
    }

    size_t bytes = (size_t)DECIM_CHUNK_FRAMES * d * 2u * src->info.out_bytes;
    src->decim_buf = (uint8_t *)malloc(bytes);

    if (src->decim_buf == NULL)
    {
        return false;
    }

    /* 对外统一暴露输出采样率，I2S / 界面都按这个值工作 */
    src->info.sample_rate /= d;

    return true;
}

/* ---------------- 抽取用的采样读写 ---------------- */

static inline int32_t pcm_sample_get(const uint8_t *p, uint8_t bytes)
{
    switch (bytes)
    {
        case 2:
            return (int16_t)(p[0] | (p[1] << 8));

        case 3:
        {
            int32_t v = (int32_t)(p[0] | (p[1] << 8) | (p[2] << 16));

            if (v & 0x800000)
            {
                v -= 0x1000000;
            }

            return v;
        }

        default:
            return (int32_t)((uint32_t)p[0] | ((uint32_t)p[1] << 8) |
                             ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24));
    }
}

static inline void pcm_sample_put(uint8_t *p, uint8_t bytes, int32_t v)
{
    switch (bytes)
    {
        case 2:
            p[0] = (uint8_t)(v & 0xFF);
            p[1] = (uint8_t)((v >> 8) & 0xFF);
            break;

        case 3:
            p[0] = (uint8_t)(v & 0xFF);
            p[1] = (uint8_t)((v >> 8) & 0xFF);
            p[2] = (uint8_t)((v >> 16) & 0xFF);
            break;

        default:
            p[0] = (uint8_t)(v & 0xFF);
            p[1] = (uint8_t)((v >> 8) & 0xFF);
            p[2] = (uint8_t)((v >> 16) & 0xFF);
            p[3] = (uint8_t)((v >> 24) & 0xFF);
            break;
    }
}

/**
 * @brief  整数抽取：把 decim 帧平均成 1 帧（简易低通，抑制抽取混叠）
 */
static void decim_mix(const uint8_t *src, uint8_t *dst, int frames, int decim, uint8_t ob)
{
    for (int i = 0; i < frames; i++)
    {
        for (int c = 0; c < 2; c++)
        {
            int64_t sum = 0;

            for (int k = 0; k < decim; k++)
            {
                const size_t idx = ((size_t)i * (size_t)decim + (size_t)k) * 2u + (size_t)c;
                sum += pcm_sample_get(src + idx * ob, ob);
            }

            pcm_sample_put(dst + ((size_t)i * 2u + (size_t)c) * ob, ob, (int32_t)(sum / decim));
        }
    }
}