/**
 * @file m4a.h
 * @brief M4A/AAC decoder interface.
 *
 * Provides decoding support for M4A and AAC-encoded audio files,
 * wrapping platform or library-specific decoding routines.
 */

#ifndef M4A_H
#define M4A_H

#ifdef __cplusplus
extern "C" {
#endif
#ifdef USE_FAAD

#include "audiotypes.h"

#include "../include/libmp4/include/libmp4.h"
#include "neaacdec.h"
#include <miniaudio.h>

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

#define MAX_CHANNELS 2
#define MAX_SAMPLES 8192 // Maximum expected frame size
#define MAX_SAMPLE_SIZE 4

typedef struct m4a_decoder {
        ma_data_source_base ds; // The m4a decoder can be used independently as a data source.
        ma_read_proc onRead;
        ma_seek_proc onSeek;
        ma_tell_proc onTell;
        void *pReadSeekTellUserData;
        ma_format format;
        FILE *mf;

        // faad2 related fields...
        NeAACDecHandle hDecoder;
        NeAACDecFrameInfo frameInfo;
        unsigned char *buffer;
        unsigned int buffer_size;
        ma_uint32 sampleSize;
        int bit_depth;
        ma_uint32 sample_rate;
        ma_uint32 channels;
        ma_uint32 avg_bit_rate;
        double duration;
        unsigned long totalFrames;

        k_m4adec_filetype file_type;

        ma_int64 file_size;

        // libmp4 fields...
        struct mp4_demux *mp4;
        struct mp4_track_info track;
        unsigned int audio_track_id;

        int32_t audio_track_index;
        uint32_t current_sample;
        uint32_t total_samples;

        uint8_t leftoverBuffer[MAX_SAMPLES *
                               MAX_CHANNELS *
                               MAX_SAMPLE_SIZE];

        ma_uint64 leftoverSampleCount;

        // For m4a_decoder_init_file
        FILE *file;

        ma_uint64 cursor;
} m4a_decoder;

#define FOUR_CHAR_INT(a, b, c, d) (((uint32_t)(a) << 24) | ((b) << 16) | ((c) << 8) | (d))

/**
 * @brief Initializes an M4A decoder using custom I/O callbacks.
 *
 * Sets up an M4A decoding context using user-provided read, seek, and
 * optional tell callbacks. Supports reading PCM frames from the
 * M4A/ALAC container format.
 *
 * @param onRead                    Callback for reading data from the source.
 * @param onSeek                    Callback for seeking within the source.
 * @param onTell                    Callback for retrieving the current position (optional).
 * @param pReadSeekTellUserData     User data passed to I/O callbacks.
 * @param p_config                  Optional decoding backend configuration.
 * @param p_allocation_callbacks    Optional allocation callbacks (currently unused).
 * @param pM4a                      Pointer to the M4A decoder to initialize.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_ARGS if required arguments are NULL.
 * @return MA_INVALID_FILE if the input data cannot be parsed as M4A.
 */
MA_API ma_result m4a_decoder_init(ma_read_proc onRead,
                                  ma_seek_proc onSeek,
                                  ma_tell_proc onTell,
                                  void *pReadSeekTellUserData,
                                  const ma_decoding_backend_config *p_config,
                                  const ma_allocation_callbacks *p_allocation_callbacks,
                                  m4a_decoder *pM4a);

/**
 * @brief Initializes an M4A decoder from a file on disk.
 *
 * Opens the specified file path and initializes decoding for the first
 * audio track (ALAC).
 *
 * @param pFilePath                 Path to the M4A file.
 * @param p_config                  Optional decoding backend configuration.
 * @param p_allocation_callbacks    Optional allocation callbacks (currently unused).
 * @param pM4a                      Pointer to the M4A decoder to initialize.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_FILE if the file cannot be opened or parsed.
 */
MA_API ma_result m4a_decoder_init_file(const char *pFilePath,
                                       const ma_decoding_backend_config *p_config,
                                       const ma_allocation_callbacks *p_allocation_callbacks,
                                       m4a_decoder *pM4a);

/**
 * @brief Uninitializes an M4A decoder and releases associated resources.
 *
 * Frees decoder state and cleans up the underlying data source.
 *
 * @param pM4a                     Pointer to the M4A decoder.
 * @param p_allocation_callbacks    Optional allocation callbacks (currently unused).
 */
MA_API void m4a_decoder_uninit(m4a_decoder *pM4a,
                               const ma_allocation_callbacks *p_allocation_callbacks);

/**
 * @brief Reads decoded PCM frames from the M4A stream.
 *
 * Decodes audio packets and writes interleaved PCM frames to the
 * output buffer. Handles internal buffering and maintains a cursor
 * for subsequent reads.
 *
 * @param pM4a         Pointer to the M4A decoder.
 * @param p_frames_out Output buffer for decoded PCM frames.
 * @param frame_count  Number of PCM frames to read.
 * @param p_frames_read Optional pointer to receive the actual number of frames read.
 *
 * @return MA_SUCCESS on success.
 * @return MA_AT_END if the end of stream is reached.
 * @return MA_INVALID_ARGS if arguments are invalid.
 * @return MA_ERROR on decoding failure.
 */
MA_API ma_result m4a_decoder_read_pcm_frames(m4a_decoder *pM4a,
                                             void *p_frames_out,
                                             ma_uint64 frame_count,
                                             ma_uint64 *p_frames_read);

/**
 * @brief Seeks to a specific PCM frame in the M4A stream.
 *
 * Performs a container-level seek and resets internal decoder state.
 *
 * @param pM4a         Pointer to the M4A decoder.
 * @param frame_index  Target PCM frame index.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_ARGS if arguments are invalid.
 * @return MA_INVALID_OPERATION if seeking fails.
 */
MA_API ma_result m4a_decoder_seek_to_pcm_frame(m4a_decoder *pM4a,
                                               ma_uint64 frame_index);

/**
 * @brief Retrieves the audio format of the decoded stream.
 *
 * Returns the sample format, channel count, sample rate, and optionally
 * fills a standard channel map.
 *
 * @param pM4a           Pointer to the M4A decoder.
 * @param p_format       Pointer to receive the sample format.
 * @param p_channels     Pointer to receive the number of channels.
 * @param p_sample_rate  Pointer to receive the sample rate.
 * @param p_channel_map  Optional buffer to receive the channel map.
 * @param channel_map_cap Capacity of the channel map buffer.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_OPERATION if decoder is invalid.
 */
MA_API ma_result m4a_decoder_get_data_format(m4a_decoder *pM4a,
                                             ma_format *p_format,
                                             ma_uint32 *p_channels,
                                             ma_uint32 *p_sample_rate,
                                             ma_channel *p_channel_map,
                                             size_t channel_map_cap);

/**
 * @brief Retrieves the current playback cursor position in PCM frames.
 *
 * @param pM4a     Pointer to the M4A decoder.
 * @param p_cursor Pointer to receive the current PCM frame index.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_ARGS if arguments are invalid.
 */
MA_API ma_result m4a_decoder_get_cursor_in_pcm_frames(m4a_decoder *pM4a,
                                                      ma_uint64 *p_cursor);

/**
 * @brief Retrieves the total length of the M4A stream in PCM frames.
 *
 * Computes and caches the total frame count based on container duration.
 *
 * @param pM4a    Pointer to the M4A decoder.
 * @param p_length Pointer to receive the total PCM frame length.
 *
 * @return MA_SUCCESS on success.
 * @return MA_INVALID_ARGS if arguments are invalid.
 */
MA_API ma_result m4a_decoder_get_length_in_pcm_frames(m4a_decoder *pM4a,
                                                      ma_uint64 *p_length);

/********************************************************************
 * Internal Data Source Adapter Functions
 * These wrap the M4A decoder API to implement the ma_data_source
 * interface.
 ********************************************************************/

/**
 * @brief Data source wrapper: retrieves the audio format from an M4A decoder.
 */
ma_result m4a_decoder_ds_get_data_format(ma_data_source *p_data_source,
                                         ma_format *p_format,
                                         ma_uint32 *p_channels,
                                         ma_uint32 *p_sample_rate,
                                         ma_channel *p_channel_map,
                                         size_t channel_map_cap);

/**
 * @brief Data source wrapper: reads PCM frames from an M4A decoder.
 */
ma_result m4a_decoder_ds_read(ma_data_source *p_data_source,
                              void *p_frames_out,
                              ma_uint64 frame_count,
                              ma_uint64 *p_frames_read);

/**
 * @brief Data source wrapper: seeks to a PCM frame in an M4A decoder.
 */
ma_result m4a_decoder_ds_seek(ma_data_source *p_data_source,
                              ma_uint64 frame_index);

/**
 * @brief Data source wrapper: retrieves the current cursor position in an M4A decoder.
 */
ma_result m4a_decoder_ds_get_cursor(ma_data_source *p_data_source,
                                    ma_uint64 *p_cursor);

/**
 * @brief Data source wrapper: retrieves the total length of the M4A decoder in PCM frames.
 */
ma_result m4a_decoder_ds_get_length(ma_data_source *p_data_source,
                                    ma_uint64 *p_length);

#if defined(MINIAUDIO_IMPLEMENTATION) || defined(MA_IMPLEMENTATION)

ma_result m4a_decoder_ds_read(ma_data_source *p_data_source, void *p_frames_out, ma_uint64 frame_count, ma_uint64 *p_frames_read)
{
        return m4a_decoder_read_pcm_frames((m4a_decoder *)p_data_source, p_frames_out, frame_count, p_frames_read);
}

ma_result m4a_decoder_ds_seek(ma_data_source *p_data_source, ma_uint64 frame_index)
{
        return m4a_decoder_seek_to_pcm_frame((m4a_decoder *)p_data_source, frame_index);
}

ma_result m4a_decoder_ds_get_data_format(ma_data_source *p_data_source, ma_format *p_format, ma_uint32 *p_channels, ma_uint32 *p_sample_rate, ma_channel *p_channel_map, size_t channel_map_cap)
{
        return m4a_decoder_get_data_format((m4a_decoder *)p_data_source, p_format, p_channels, p_sample_rate, p_channel_map, channel_map_cap);
}

ma_result m4a_decoder_ds_get_cursor(ma_data_source *p_data_source, ma_uint64 *p_cursor)
{
        return m4a_decoder_get_cursor_in_pcm_frames((m4a_decoder *)p_data_source, p_cursor);
}

ma_result m4a_decoder_ds_get_length(ma_data_source *p_data_source, ma_uint64 *p_length)
{
        return m4a_decoder_get_length_in_pcm_frames((m4a_decoder *)p_data_source, p_length);
}

ma_data_source_vtable g_m4a_decoder_ds_vtable =
    {
        m4a_decoder_ds_read,
        m4a_decoder_ds_seek,
        m4a_decoder_ds_get_data_format,
        m4a_decoder_ds_get_cursor,
        m4a_decoder_ds_get_length,
        NULL,
        (ma_uint64)0};

static ma_result file_on_seek(void *pUserData, ma_int64 offset, ma_seek_origin origin)
{
        FILE *fp = (FILE *)pUserData;
        int whence = (origin == ma_seek_origin_start) ? SEEK_SET : SEEK_CUR;
        if (fseeko(fp, offset, whence) != 0) {
                return MA_ERROR;
        }
        return MA_SUCCESS;
}

static ma_result m4a_decoder_init_internal(const ma_decoding_backend_config *p_config, m4a_decoder *pM4a)
{
        if (pM4a == NULL) {
                return MA_INVALID_ARGS;
        }

        MA_ZERO_OBJECT(pM4a);
        pM4a->format = ma_format_f32;

        if (p_config != NULL && (p_config->preferredFormat == ma_format_f32 || p_config->preferredFormat == ma_format_s16)) {
                pM4a->format = p_config->preferredFormat;
        }

        ma_data_source_config dataSourceConfig = ma_data_source_config_init();

        dataSourceConfig.vtable = &g_m4a_decoder_ds_vtable;

        ma_result result = ma_data_source_init(&dataSourceConfig, &pM4a->ds);
        if (result != MA_SUCCESS) {
                return result;
        }

        return MA_SUCCESS;
}

// Note: This isn't used by kew and is untested
MA_API ma_result m4a_decoder_init(
    ma_read_proc onRead,
    ma_seek_proc onSeek,
    ma_tell_proc onTell,
    void *pReadSeekTellUserData,
    const ma_decoding_backend_config *p_config,
    const ma_allocation_callbacks *p_allocation_callbacks,
    m4a_decoder *pM4a)
{
        (void)onRead;
        (void)onSeek;
        (void)onTell;
        (void)pReadSeekTellUserData;
        (void)p_config;
        (void)p_allocation_callbacks;
        (void)pM4a;

        /*
        if (pM4a == NULL || onRead == NULL || onSeek == NULL || onTell == NULL) {
                return MA_INVALID_ARGS;
        }

        ma_result result = m4a_decoder_init_internal(p_config, pM4a);
        if (result != MA_SUCCESS) {
                return result;
        }

        // Store the custom read, seek, and tell functions
        pM4a->pReadSeekTellUserData = pReadSeekTellUserData;

        // Get the size of the data source
        ma_int64 currentPos = 0;
        if (pM4a->onTell(pM4a->pReadSeekTellUserData, &currentPos) != MA_SUCCESS) {
                return MA_ERROR;
        }

        if (pM4a->onSeek(pM4a->pReadSeekTellUserData, 0, ma_seek_origin_end) != MA_SUCCESS) {
                return MA_ERROR;
        }

        ma_int64 file_size = 0;
        if (pM4a->onTell(pM4a->pReadSeekTellUserData, &file_size) != MA_SUCCESS) {
                return MA_ERROR;
        }

        // Seek back to original position
        if (pM4a->onSeek(pM4a->pReadSeekTellUserData, currentPos, ma_seek_origin_start) != MA_SUCCESS) {
                return MA_ERROR;
        }

        // Initialize mp4lib
        if (mp4_demux_open(filename, &pM4a->mp4) < 0)
                return MA_ERROR;

        pM4a->buffer = NULL;

        // Find the audio track
        pM4a->audio_track_index = -1;
        for (unsigned int i = 0; i < pM4a->mp4->track_count; i++) {
                MP4D_track_t *track = &pM4a->mp4.track[i];
                if (track->handler_type == MP4D_HANDLER_TYPE_SOUN) {
                        pM4a->audio_track_index = i;
                        pM4a->track = track;
                        break;
                }
        }

        if (pM4a->audio_track_index == -1) {
                // No audio track found
                mp4_demux_close(pM4a->mp4);
                return MA_ERROR;
        }

        pM4a->current_sample = 0;
        pM4a->total_samples = pM4a->track->sample_count;

        // Initialize faad2 decoder
        pM4a->hDecoder = NeAACDecOpen();

        // Extract the decoder configuration
        const uint8_t *decoder_config = pM4a->track->dsi;
        uint32_t decoder_config_len = pM4a->track->dsi_bytes;

        unsigned long sample_rate;
        unsigned char channels;

        if (NeAACDecInit2(pM4a->hDecoder, (unsigned char *)decoder_config, decoder_config_len, &sample_rate, &channels) < 0) {
                // Error initializing decoder
                NeAACDecClose(pM4a->hDecoder);
                mp4_demux_close(pM4a->mp4);
                return MA_ERROR;
        }

        // Configure output format
        NeAACDecConfigurationPtr config = NeAACDecGetCurrentConfiguration(pM4a->hDecoder);
        if (pM4a->format == ma_format_s16) {
                config->outputFormat = FAAD_FMT_16BIT;
                pM4a->sampleSize = sizeof(int16_t);
                pM4a->bit_depth = 16;
        } else if (pM4a->format == ma_format_f32) {
                config->outputFormat = FAAD_FMT_FLOAT;
                pM4a->sampleSize = sizeof(float);
                pM4a->bit_depth = 32;
        } else {
                // Unsupported format
                NeAACDecClose(pM4a->hDecoder);
                mp4_demux_close(pM4a->mp4);
                return MA_ERROR;
        }
        NeAACDecSetConfiguration(pM4a->hDecoder, config);

        // Initialize other fields
        pM4a->leftoverSampleCount = 0;
        pM4a->cursor = 0;
        */

        fprintf(stderr, " m4a_decoder_init: Not Implemented");
        fflush(stderr);
        exit(0);

        return MA_SUCCESS;
}

double calculate_aac_duration(FILE *fp, unsigned long sample_rate, unsigned long *totalFrames)
{
        if (fp == NULL || sample_rate == 0 || totalFrames == NULL) {
                return -1.0;
        }

        unsigned char buffer[7];
        unsigned long file_size = 0;
        *totalFrames = 0;

        // Get file size
        fseek(fp, 0, SEEK_END);
        file_size = ftell(fp);
        fseek(fp, 0, SEEK_SET);

        // Loop to count frames
        while (ftell(fp) < (long)(file_size - 7)) // Ensure at least an ADTS header remains
        {
                // Read header
                if (fread(buffer, 1, 7, fp) < 7)
                        break;

                // Extract frame size
                unsigned int frameSize = ((buffer[3] & 0x03) << 11) | ((buffer[4] & 0xFF) << 3) | ((buffer[5] & 0xE0) >> 5);

                if (frameSize <= 7)
                        break;

                // Skip to next frame
                fseek(fp, frameSize - 7, SEEK_CUR);

                (*totalFrames)++;
        }

        // Compute duration using: duration = (totalFrames * 1024) / sample_rate
        double duration = (double)(*totalFrames * 1024) / sample_rate;

        fseek(fp, 0, SEEK_SET);

        return duration;
}

uint32_t read_u32be(FILE *fp)
{
        unsigned char b[4];
        if (fread(b, 1, 4, fp) != 4)
                return 0;
        return ((uint32_t)b[0] << 24) | ((uint32_t)b[1] << 16) | ((uint32_t)b[2] << 8) | ((uint32_t)b[3]);
}

int find_atom(FILE *fp, uint32_t atom_name, long max_search_length, uint32_t *atom_size_out)
{
        long start_pos = ftell(fp);
        while ((ftell(fp) - start_pos) < max_search_length) {
                unsigned char header[8];
                if (fread(header, 1, 8, fp) != 8)
                        return 0;

                uint32_t atom_size = (header[0] << 24) | (header[1] << 16) | (header[2] << 8) | header[3];
                uint32_t atom_type = (header[4] << 24) | (header[5] << 16) | (header[6] << 8) | header[7];

                if (atom_size < 8)
                        return 0; // Invalid atom size

                if (atom_type == atom_name) {
                        if (atom_size_out)
                                *atom_size_out = atom_size - 8;

                        return 1; // Found
                }

                if (fseek(fp, atom_size - 8, SEEK_CUR) != 0)
                        return 0;
        }
        return 0; // Not found
}

int is_alac(FILE *fp, uint8_t *dsi_out, size_t *dsi_size_out)
{
        fseek(fp, 0, SEEK_SET);
        uint32_t atom_size;

        if (!find_atom(fp, FOUR_CHAR_INT('m', 'o', 'o', 'v'), 0x7FFFFFFF, &atom_size))
                return 0;
        if (!find_atom(fp, FOUR_CHAR_INT('t', 'r', 'a', 'k'), atom_size, &atom_size))
                return 0;
        if (!find_atom(fp, FOUR_CHAR_INT('m', 'd', 'i', 'a'), atom_size, &atom_size))
                return 0;
        if (!find_atom(fp, FOUR_CHAR_INT('m', 'i', 'n', 'f'), atom_size, &atom_size))
                return 0;
        if (!find_atom(fp, FOUR_CHAR_INT('s', 't', 'b', 'l'), atom_size, &atom_size))
                return 0;
        if (!find_atom(fp, FOUR_CHAR_INT('s', 't', 's', 'd'), atom_size, &atom_size))
                return 0;

        fseek(fp, 8, SEEK_CUR); // Skip stsd header (version+entry)

        read_u32be(fp); // uint32_t sample_entry_size
        uint32_t sample_entry_fourcc = read_u32be(fp);
        if (sample_entry_fourcc != FOUR_CHAR_INT('a', 'l', 'a', 'c'))
                return 0;

        fseek(fp, 28, SEEK_CUR); // Skip audio sample entry fields

        uint32_t config_atom_size = read_u32be(fp);
        uint32_t config_atom_fourcc = read_u32be(fp);
        if (config_atom_fourcc != FOUR_CHAR_INT('a', 'l', 'a', 'c'))
                return 0;

        fseek(fp, 4, SEEK_CUR); // Skip 1-byte version and 3-byte flags (4 bytes total)!

        uint32_t alac_dsi_size = config_atom_size - 12; // size(4)+fourcc(4)+version/flags(4) total=12 bytes overhead
        if (alac_dsi_size < 24 || alac_dsi_size > 64)
                return 0; // Sanity check

        if (fread(dsi_out, 1, alac_dsi_size, fp) != alac_dsi_size)
                return 0;
        *dsi_size_out = alac_dsi_size;

        return 1;
}

MA_API ma_result m4a_decoder_init_file(
    const char *pFilePath,
    const ma_decoding_backend_config *p_config,
    const ma_allocation_callbacks *p_allocation_callbacks,
    m4a_decoder *pM4a)
{
        (void)p_allocation_callbacks;

        if (pFilePath == NULL || pM4a == NULL) {
                return MA_INVALID_ARGS;
        }

        ma_result result = m4a_decoder_init_internal(p_config, pM4a);
        if (result != MA_SUCCESS) {
                return result;
        }

        FILE *fp = fopen(pFilePath, "rb");
        if (fp == NULL) {
                return MA_INVALID_FILE;
        }

        // Get the file size
        if (fseeko(fp, 0, SEEK_END) != 0) {
                fclose(fp);
                return MA_ERROR;
        }

        pM4a->file_size = ftello(fp);
        if (pM4a->file_size < 0) {
                fclose(fp);
                return MA_ERROR;
        }

        if (fseeko(fp, 0, SEEK_SET) != 0) {
                fclose(fp);
                return MA_ERROR;
        }

        // Store the FILE pointer in the decoder struct
        pM4a->file = fp;

        // Try to detect the file format (ADTS, MP4, LATM, etc.)
        unsigned char buffer[7];
        size_t bytes_read = fread(buffer, 1, sizeof(buffer), fp);

        // Check for ADTS header
        if (bytes_read >= 7 && buffer[0] == 0xFF && (buffer[1] & 0xF0) == 0xF0) {
                pM4a->file_type = k_rawAAC;
        } else {
                // Check if it's an MP4 file (using MP4D_open or similar)
                if (mp4_demux_open(pFilePath, &pM4a->mp4) < 0) {
                        fclose(fp);
                        return MA_ERROR;
                } else {
                        pM4a->file_type = k_unknown; // It's an MP4 container
                }
        }

        pM4a->buffer = NULL;

        if (pM4a->file_type == k_rawAAC) {
                // Raw AAC handling

                // Extract the frame size from the ADTS header
                unsigned int frameSize = ((buffer[3] & 0x03) << 11) | ((buffer[4] & 0xFF) << 3) | ((buffer[5] & 0xE0) >> 5);

                if (frameSize <= 7) {
                        fclose(fp);
                        return MA_ERROR; // Invalid frame size
                }

                unsigned char *frameData = malloc(frameSize);
                if (frameData == NULL) {
                        fclose(fp);
                        return MA_ERROR; // Memory allocation failed
                }

                // The first 7 bytes are already in the buffer, so copy them to frameData
                memcpy(frameData, buffer, 7);

                // Read the rest of the frame (audio data)
                size_t remainingBytes = frameSize - 7;
                size_t additionalBytesRead = fread(frameData + 7, 1, remainingBytes, fp);
                if (additionalBytesRead < remainingBytes) {
                        free(frameData);
                        fclose(fp);
                        return MA_ERROR; // Failed to read the full frame
                }

                // Allocate decoder config
                unsigned char *decoder_config = malloc(2);
                if (decoder_config == NULL) {
                        return MA_OUT_OF_MEMORY;
                }

                unsigned long decoder_config_size = 2;

                decoder_config[0] = ((buffer[2] & 0xC0) >> 6) + 1; // Object type
                decoder_config[0] |= ((buffer[2] & 0x3C) >> 2) << 2;

                decoder_config[1] = ((buffer[2] & 0x07) << 1) | ((buffer[3] & 0x80) >> 7); // Channels and sample_rate
                decoder_config[1] <<= 4;                                                   // Shift to upper 4 bits

                unsigned char objectType = decoder_config[0];
                if (objectType == 5 || objectType >= 29) {
                        fprintf(stderr, "File is encoded with HE-AAC which is not supported.\n");
                        free(frameData);
                        free(decoder_config);
                        fclose(fp);
                        return MA_ERROR;
                }

                unsigned long sample_rate = 0;
                unsigned char channels = 0;

                pM4a->hDecoder = NeAACDecOpen();

                int initResult = NeAACDecInit2(pM4a->hDecoder, (unsigned char *)decoder_config, decoder_config_size, &sample_rate, &channels);
                if (initResult < 0) {
                        fprintf(stderr, "Error initializing decoder. Code: %d\n", initResult);
                        free(frameData);
                        free(decoder_config);
                        NeAACDecClose(pM4a->hDecoder);
                        fclose(fp);
                        return MA_ERROR;
                }

                free(decoder_config);

                // Check if the sample_rate and channels are correctly initialized
                if (sample_rate == 0 || channels == 0) {
                        fprintf(stderr, "Error: Invalid sample rate or channel count.\n");
                        free(frameData);
                        NeAACDecClose(pM4a->hDecoder);
                        fclose(fp);
                        return MA_ERROR;
                }

                pM4a->sample_rate = (ma_uint32)sample_rate;
                pM4a->channels = (ma_uint32)channels;

                // Clean up the frame data after processing
                free(frameData);

                // Configure output format
                NeAACDecConfigurationPtr config_ptr = NeAACDecGetCurrentConfiguration(pM4a->hDecoder);
                if (pM4a->format == ma_format_s16) {
                        config_ptr->outputFormat = FAAD_FMT_16BIT;
                        pM4a->sampleSize = sizeof(int16_t);
                        pM4a->bit_depth = 16;
                } else if (pM4a->format == ma_format_f32) {
                        config_ptr->outputFormat = FAAD_FMT_FLOAT;
                        pM4a->sampleSize = sizeof(float);
                        pM4a->bit_depth = 32;
                } else {
                        // Unsupported format
                        NeAACDecClose(pM4a->hDecoder);
                        fclose(fp);
                        return MA_ERROR;
                }
                NeAACDecSetConfiguration(pM4a->hDecoder, config_ptr);

                // Initialize other fields
                pM4a->leftoverSampleCount = 0;
                pM4a->cursor = 0;

                fseek(pM4a->file, 0, SEEK_SET);

                return MA_SUCCESS;
        } else {
                // Find the audio track

                int trackCount = mp4_demux_get_track_count(pM4a->mp4);

                pM4a->audio_track_index = -1;

                for (int i = 0; i < trackCount; i++) {

                        struct mp4_track_info info;

                        if (mp4_demux_get_track_info(pM4a->mp4, i, &info) < 0)
                                continue;

                        if (info.type == MP4_TRACK_TYPE_AUDIO) {

                                pM4a->audio_track_index = i;
                                pM4a->audio_track_id = info.id;
                                pM4a->track = info;

                                pM4a->buffer = malloc(info.sample_max_size);
                                pM4a->buffer_size = info.sample_max_size;

                                break;
                        }
                }

                // M4A (MP4-wrapped AAC) handling
                if (pM4a->audio_track_index == -1) {
                        // No audio track found
                        mp4_demux_close(pM4a->mp4);
                        fclose(fp);
                        return MA_ERROR;
                }

                pM4a->current_sample = 0;
                pM4a->total_samples = pM4a->track.sample_count;

                uint64_t duration_us =
                    mp4_sample_time_to_usec(pM4a->track.duration,
                                            pM4a->track.timescale);

                if (duration_us > 0) {
                        uint64_t bps = (pM4a->file_size * 8 * 1000000ULL) / duration_us;
                        pM4a->avg_bit_rate = (bps + 500) / 1000; // round to nearest kbps
                }

                uint8_t alac_dsi[32];
                size_t alac_dsi_size;

                long original_position = ftell(pM4a->file);

                if (is_alac(fp, alac_dsi, &alac_dsi_size)) {
                        // This is an alac file and is currently unsupported.
                        fprintf(stderr, "M4a files that use the ALAC encoder are not supported.");
                        return MA_ERROR;
                } else // AAC
                {
                        pM4a->file_type = k_aac;

                        fseek(pM4a->file, original_position, SEEK_SET);

                        // Initialize faad2 decoder
                        pM4a->hDecoder = NeAACDecOpen();

                        // Extract the decoder configuration
                        uint8_t *decoder_config = NULL;
                        unsigned int decoder_config_len = 0;

                        if (mp4_demux_get_track_audio_specific_config(
                                pM4a->mp4,
                                pM4a->track.id,
                                &decoder_config,
                                &decoder_config_len) != 0) {
                                return MA_ERROR;
                        }

                        unsigned long sample_rate;
                        unsigned char channels;

                        if (decoder_config_len >= 2) {
                                uint8_t object_type = (decoder_config[0] >> 3) & 0x1F;

                                if (object_type == 5 || object_type == 29) {
                                        fprintf(stderr, "Unsupported AAC object type: (HE-AAC or PS)");
                                        return MA_ERROR;
                                }
                        }

                        if (NeAACDecInit2(pM4a->hDecoder, (unsigned char *)decoder_config, decoder_config_len, &sample_rate, &channels) < 0) {
                                // Error initializing decoder
                                NeAACDecClose(pM4a->hDecoder);
                                mp4_demux_close(pM4a->mp4);
                                fclose(fp);
                                return MA_ERROR;
                        }

                        pM4a->sample_rate = (ma_uint32)sample_rate;
                        pM4a->channels = (ma_uint32)channels;

                        // Configure output format
                        NeAACDecConfigurationPtr config_ptr = NeAACDecGetCurrentConfiguration(pM4a->hDecoder);
                        if (pM4a->format == ma_format_s16) {
                                config_ptr->outputFormat = FAAD_FMT_16BIT;
                                pM4a->sampleSize = sizeof(int16_t);
                                pM4a->bit_depth = 16;
                        } else if (pM4a->format == ma_format_f32) {
                                config_ptr->outputFormat = FAAD_FMT_FLOAT;
                                pM4a->sampleSize = sizeof(float);
                                pM4a->bit_depth = 32;
                        } else {
                                // Unsupported format
                                NeAACDecClose(pM4a->hDecoder);
                                mp4_demux_close(pM4a->mp4);
                                fclose(fp);
                                return MA_ERROR;
                        }
                        NeAACDecSetConfiguration(pM4a->hDecoder, config_ptr);

                        // Initialize other fields
                        pM4a->leftoverSampleCount = 0;
                        pM4a->cursor = 0;
                        fseek(pM4a->file, 0, SEEK_SET);

                        return MA_SUCCESS;
                }
        }
}

MA_API void m4a_decoder_uninit(m4a_decoder *pM4a, const ma_allocation_callbacks *p_allocation_callbacks)
{
        (void)p_allocation_callbacks;

        if (pM4a == NULL) {
                return;
        }

        if (pM4a->hDecoder) {
                NeAACDecClose(pM4a->hDecoder);
                pM4a->hDecoder = NULL;
        }

        if (pM4a->file_type != k_rawAAC) {
                mp4_demux_close(pM4a->mp4);
        }

        if (pM4a->buffer) {
                free(pM4a->buffer);
                pM4a->buffer = NULL;
                pM4a->buffer_size = 0;
        }

        if (pM4a->file) {
                fclose(pM4a->file);
                pM4a->file = NULL;
        }
}

MA_API ma_result m4a_decoder_read_pcm_frames(
    m4a_decoder *pM4a,
    void *p_frames_out,
    ma_uint64 frame_count,
    ma_uint64 *p_frames_read)
{
        if (pM4a == NULL || p_frames_out == NULL || frame_count == 0) {
                return MA_INVALID_ARGS;
        }

        ma_result result = MA_SUCCESS;
        ma_uint32 channels = pM4a->channels;
        ma_uint32 sampleSize = pM4a->sampleSize;
        ma_uint64 totalFramesProcessed = 0;

        // Handle any leftover samples from previous call using the global/static leftover buffer
        if (pM4a->leftoverSampleCount > 0) {
                ma_uint64 leftoverToProcess = (pM4a->leftoverSampleCount < frame_count) ? pM4a->leftoverSampleCount : frame_count;
                ma_uint64 leftoverBytes = leftoverToProcess * channels * sampleSize;

                memcpy(p_frames_out, pM4a->leftoverBuffer, leftoverBytes);
                totalFramesProcessed += leftoverToProcess;

                // Shift the leftover buffer
                ma_uint64 samplesLeft = pM4a->leftoverSampleCount - leftoverToProcess;
                if (samplesLeft > 0) {
                        memmove(pM4a->leftoverBuffer, pM4a->leftoverBuffer + leftoverBytes, samplesLeft * channels * sampleSize);
                }
                pM4a->leftoverSampleCount = samplesLeft;
        }

        while (totalFramesProcessed < frame_count) {
                if (pM4a->file_type == k_rawAAC) {
                        unsigned int headerSize = 7;
                        uint8_t buffer[headerSize];

                        if (fread(buffer, 1, headerSize, pM4a->file) != headerSize) {
                                result = MA_ERROR;
                                break;
                        }

                        unsigned int frame_bytes = ((buffer[3] & 0x03) << 11) | ((buffer[4] & 0xFF) << 3) | ((buffer[5] & 0xE0) >> 5);

                        if (frame_bytes < headerSize || frame_bytes > 8192) {
                                result = MA_ERROR;
                                break;
                        }

                        // Allocate memory for the frame
                        unsigned char *sample_data = (unsigned char *)malloc(frame_bytes);
                        if (!sample_data) {
                                result = MA_OUT_OF_MEMORY;
                                break;
                        }

                        // Copy the header to the sample_data buffer
                        memcpy(sample_data, buffer, headerSize);

                        // Read the rest of the frame (audio data)
                        size_t remaining_bytes = frame_bytes - headerSize;
                        size_t additionalBytesRead = fread(sample_data + headerSize, 1, remaining_bytes, pM4a->file);

                        if (additionalBytesRead < remaining_bytes) {
                                free(sample_data);
                                result = MA_ERROR;
                                break; // Failed to read full frame
                        }

                        pM4a->current_sample++;

                        // Decode the AAC frame using faad2
                        void *decodedData = NeAACDecDecode(pM4a->hDecoder, &(pM4a->frameInfo), sample_data + 7, frame_bytes - 7);
                        free(sample_data);

                        if (pM4a->frameInfo.error > 0) {
                                // Error in decoding, skip to the next frame.
                                continue;
                        }

                        // Remove support for HE-AAC components (SBR or PS)
                        if (pM4a->frameInfo.sbr || pM4a->frameInfo.ps) {
                                // File is encoded with HE-AAC which is not supported
                                return MA_ERROR;
                        }

                        unsigned long samplesDecoded = pM4a->frameInfo.samples; // Total samples decoded (channels * frames)
                        ma_uint64 framesDecoded = samplesDecoded / channels;

                        // Calculate how many frames we can process in this call
                        ma_uint64 framesNeeded = frame_count - totalFramesProcessed;
                        ma_uint64 frames_to_copy = (framesDecoded < framesNeeded) ? framesDecoded : framesNeeded;
                        ma_uint64 bytesToCopy = frames_to_copy * channels * sampleSize;

                        memcpy((uint8_t *)p_frames_out + totalFramesProcessed * channels * sampleSize, decodedData, bytesToCopy);
                        totalFramesProcessed += frames_to_copy;

                        // Handle leftover frames using the global/static leftover buffer
                        if (frames_to_copy < framesDecoded) {
                                // There are leftover frames
                                pM4a->leftoverSampleCount = framesDecoded - frames_to_copy;
                                ma_uint64 leftoverBytes = pM4a->leftoverSampleCount * channels * sampleSize;

                                if (leftoverBytes > sizeof(pM4a->leftoverBuffer)) {
                                        // Safety check to avoid overflow in the buffer.
                                        pM4a->leftoverSampleCount = sizeof(pM4a->leftoverBuffer) / (channels * sampleSize);
                                        leftoverBytes = pM4a->leftoverSampleCount * channels * sampleSize;
                                }

                                memcpy(pM4a->leftoverBuffer, (uint8_t *)decodedData + bytesToCopy, leftoverBytes);
                        } else {
                                pM4a->leftoverSampleCount = 0;
                        }
                } else {
                        if (pM4a->current_sample >= pM4a->total_samples) {
                                result = MA_AT_END;
                                break; // No more samples
                        }

                        unsigned int frame_bytes = 0;

                        // Get the sample offset and size
                        struct mp4_track_sample sample;

                        int ret =
                            mp4_demux_get_track_sample(
                                pM4a->mp4,
                                pM4a->audio_track_id,
                                pM4a->current_sample,
                                pM4a->buffer,
                                pM4a->buffer_size,
                                NULL,
                                0,
                                &sample);

                        if (ret < 0) {
                                // Error getting sample info
                                result = MA_ERROR;
                                break;
                        }

                        frame_bytes = sample.size;

                        pM4a->current_sample++;

                        void *decodedData =
                            NeAACDecDecode(
                                pM4a->hDecoder,
                                &pM4a->frameInfo,
                                pM4a->buffer,
                                frame_bytes);

                        if (pM4a->frameInfo.error > 0) {
                                fprintf(stderr,
                                        "Decoding Error %d: %s\n",
                                        pM4a->frameInfo.error,
                                        NeAACDecGetErrorMessage(pM4a->frameInfo.error));
                                continue;
                        }

                        // Remove support for HE-AAC components (SBR or PS)
                        if (pM4a->frameInfo.sbr || pM4a->frameInfo.ps) {
                                // HE - AAC detected(either SBR or PS is present), skip processing continue;
                                continue;
                        }

                        unsigned long samplesDecoded = pM4a->frameInfo.samples; // Total samples decoded (channels * frames)
                        ma_uint64 framesDecoded = samplesDecoded / channels;

                        // Calculate how many frames we can process in this call
                        ma_uint64 framesNeeded = frame_count - totalFramesProcessed;
                        ma_uint64 frames_to_copy = (framesDecoded < framesNeeded) ? framesDecoded : framesNeeded;
                        ma_uint64 bytesToCopy = frames_to_copy * channels * sampleSize;

                        memcpy((uint8_t *)p_frames_out + totalFramesProcessed * channels * sampleSize, decodedData, bytesToCopy);
                        totalFramesProcessed += frames_to_copy;

                        // Handle leftover frames using the global/static leftover buffer
                        if (frames_to_copy < framesDecoded) {
                                // There are leftover frames
                                pM4a->leftoverSampleCount = framesDecoded - frames_to_copy;
                                ma_uint64 leftoverBytes = pM4a->leftoverSampleCount * channels * sampleSize;

                                if (leftoverBytes > sizeof(pM4a->leftoverBuffer)) {
                                        // Safety check to avoid overflow in the buffer.
                                        pM4a->leftoverSampleCount = sizeof(pM4a->leftoverBuffer) / (channels * sampleSize);
                                        leftoverBytes = pM4a->leftoverSampleCount * channels * sampleSize;
                                }

                                memcpy(pM4a->leftoverBuffer, (uint8_t *)decodedData + bytesToCopy, leftoverBytes);
                        } else {
                                pM4a->leftoverSampleCount = 0;
                        }
                }
        }

        pM4a->cursor += totalFramesProcessed;

        if (p_frames_read != NULL) {
                *p_frames_read = totalFramesProcessed;
        }

        if (totalFramesProcessed > 0 && result != MA_AT_END)
                return MA_SUCCESS;

        return result;
}

MA_API ma_result m4a_decoder_seek_to_pcm_frame(m4a_decoder *pM4a, ma_uint64 frame_index)
{
        if (pM4a == NULL)
                return MA_INVALID_ARGS;

        if (frame_index >= pM4a->total_samples * 1024)
                return MA_INVALID_ARGS;

        pM4a->current_sample = (ma_uint32)(frame_index / 1024);

        if (pM4a->file_type == k_rawAAC) {
                return MA_ERROR;
        } else if (pM4a->file_type == k_ALAC) {
                unsigned int frame_bytes = 0;

                // Get the sample offset and size
                struct mp4_track_sample sample;

                int ret =
                    mp4_demux_get_track_sample(
                        pM4a->mp4,
                        pM4a->audio_track_id,
                        1,
                        pM4a->buffer,
                        pM4a->buffer_size,
                        NULL,
                        0,
                        &sample);

                if (ret < 0)
                        // Error getting sample info
                        return MA_ERROR;

                frame_bytes = sample.size;
                ma_int64 sample_offset = sample.offset;

                if (sample_offset < 0 || frame_bytes == 0) {
                        return MA_ERROR;
                }

                if (file_on_seek(pM4a->file, sample_offset, ma_seek_origin_start) != MA_SUCCESS) {
                        return MA_ERROR;
                }

                pM4a->leftoverSampleCount = 0;

                uint64_t actual_pcm_frame =
                    (sample.dts * (uint64_t)pM4a->sample_rate) / pM4a->track.timescale;

                pM4a->current_sample = (ma_uint32)(actual_pcm_frame / 1024);
                pM4a->cursor = actual_pcm_frame;

                return MA_SUCCESS;
        } else {
                unsigned int frame_bytes = 0;
                struct mp4_track_sample sample;

                // frame_index -> PCM frame count -> microseconds
                // time_us = frames / sample_rate * 1,000,000
                uint64_t time_us =
                    ((uint64_t)pM4a->current_sample * 1024ULL * 1000000ULL) / pM4a->sample_rate;

                int seek_ret = mp4_demux_seek(
                    pM4a->mp4,
                    time_us,
                    MP4_SEEK_METHOD_PREVIOUS_SYNC);

                if (seek_ret < 0)
                        return MA_ERROR;

                int ret =
                    mp4_demux_get_track_sample(
                        pM4a->mp4,
                        pM4a->audio_track_id,
                        pM4a->current_sample,
                        NULL,
                        0,
                        NULL,
                        0,
                        &sample);

                if (ret < 0)
                        return MA_ERROR;

                frame_bytes = sample.size;
                ma_int64 sample_offset = sample.offset;

                if (sample_offset < 0 || frame_bytes == 0) {
                        return MA_ERROR;
                }

                uint64_t actual_pcm_frame =
                    (sample.dts * (uint64_t)pM4a->sample_rate) / pM4a->track.timescale;

                pM4a->current_sample = (ma_uint32)(actual_pcm_frame / 1024);

                NeAACDecPostSeekReset(pM4a->hDecoder, (long)pM4a->current_sample);

                pM4a->leftoverSampleCount = 0;
                pM4a->cursor = actual_pcm_frame;

                return MA_SUCCESS;
        }
}

MA_API ma_result m4a_decoder_get_data_format(
    m4a_decoder *pM4a,
    ma_format *p_format,
    ma_uint32 *p_channels,
    ma_uint32 *p_sample_rate,
    ma_channel *p_channel_map,
    size_t channel_map_cap)
{
        // Initialize output variables
        if (p_format != NULL) {
                *p_format = ma_format_unknown;
        }
        if (p_channels != NULL) {
                *p_channels = 0;
        }
        if (p_sample_rate != NULL) {
                *p_sample_rate = 0;
        }
        if (p_channel_map != NULL) {
                MA_ZERO_MEMORY(p_channel_map, sizeof(*p_channel_map) * channel_map_cap);
        }

        if (pM4a == NULL) {
                return MA_INVALID_OPERATION;
        }

        if (p_format != NULL) {
                *p_format = pM4a->format;
        }

        if (p_channels != NULL) {
                *p_channels = pM4a->channels;
        }

        if (p_sample_rate != NULL) {
                *p_sample_rate = pM4a->sample_rate;
        }

        // Set a standard channel map if requested
        if (p_channel_map != NULL) {
                ma_channel_map_init_standard(ma_standard_channel_map_microsoft, p_channel_map, channel_map_cap, *p_channels);
        }

        return MA_SUCCESS;
}

MA_API ma_result m4a_decoder_get_cursor_in_pcm_frames(m4a_decoder *pM4a, ma_uint64 *p_cursor)
{
        if (p_cursor == NULL) {
                return MA_INVALID_ARGS;
        }

        *p_cursor = 0; /* Safety. */

        if (pM4a == NULL) {
                return MA_INVALID_ARGS;
        }

        *p_cursor = pM4a->cursor;

        return MA_SUCCESS;
}

MA_API ma_result m4a_decoder_get_length_in_pcm_frames(m4a_decoder *pM4a,
                                                      ma_uint64 *p_length)
{
        if (p_length == NULL) {
                return MA_INVALID_ARGS;
        }

        *p_length = 0;

        if (pM4a == NULL) {
                return MA_INVALID_ARGS;
        }

        if (pM4a->sample_rate == 0 || pM4a->track.timescale == 0) {
                return MA_ERROR;
        }

        *p_length =
            ((ma_uint64)pM4a->track.duration * pM4a->sample_rate) /
            pM4a->track.timescale;

        return MA_SUCCESS;
}
#endif
#endif
#ifdef __cplusplus
}
#endif
#endif
