#include "audio_compressor.h"
#include "debug_config.h"
#include "mem.h"
#include <math.h>
#include <string.h>
#include <stdlib.h>

#define BLOCK_LEVELS_CHUNK 64
#define PENDING_CHUNK       32

struct audio_compressor {
    audio_compressor_config_t cfg;
    int enabled;

    int block_samples;
    int lookback_blocks;
    int lookahead_blocks;
    float rise_factor_per_block;
    float max_gain;
    float gain_smoothed;

    float*   block_levels;
    size_t   block_levels_count;
    size_t   block_levels_cap;

    struct pending_block {
        int16_t* samples;
        size_t   count;
        size_t   level_index;
    }*       pending;
    size_t   pending_count;
    size_t   pending_cap;

    int16_t* accumulator;
    size_t   accum_count;

    int16_t* output;
    size_t   output_size;
    size_t   output_cap;

    int dbg_counter;
};

struct audio_compressor* audio_compressor_create(void) {
    struct audio_compressor* ac = u_calloc(1, sizeof(*ac));
    if (!ac) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "audio_compressor_create: OOM");
        return NULL;
    }
    ac->enabled = 1;
    ac->gain_smoothed = 1.0f;
    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "audio_compressor_create: ok");
    return ac;
}

void audio_compressor_destroy(struct audio_compressor* ac) {
    if (!ac) return;
    for (size_t i = 0; i < ac->pending_count; i++)
        u_free(ac->pending[i].samples);
    u_free(ac->pending);
    u_free(ac->block_levels);
    u_free(ac->accumulator);
    u_free(ac->output);
    u_free(ac);
    DEBUG_DEBUG(DEBUG_CATEGORY_GENERAL, "audio_compressor_destroy");
}

void audio_compressor_configure(struct audio_compressor* ac, const audio_compressor_config_t* cfg) {
    if (!ac || !cfg) return;
    ac->cfg = *cfg;
    if (ac->cfg.sample_rate <= 0) ac->cfg.sample_rate = 48000;
    if (ac->cfg.channels <= 0)    ac->cfg.channels = 1;
    if (ac->cfg.block_duration_ms <= 0) ac->cfg.block_duration_ms = 20;
    if (ac->cfg.lookback_ms <= 0)  ac->cfg.lookback_ms = 200;
    if (ac->cfg.lookahead_ms <= 0) ac->cfg.lookahead_ms = 100;
    if (ac->cfg.target_level <= 0.0f) ac->cfg.target_level = 0.25f;

    ac->block_samples = ac->cfg.sample_rate * ac->cfg.block_duration_ms / 1000 * ac->cfg.channels;
    ac->lookback_blocks  = ac->cfg.lookback_ms / ac->cfg.block_duration_ms;
    ac->lookahead_blocks = ac->cfg.lookahead_ms / ac->cfg.block_duration_ms;
    ac->rise_factor_per_block = powf(ac->cfg.rise_rate_per_sec, (float)ac->cfg.block_duration_ms / 1000.0f);
    ac->max_gain = powf(10.0f, ac->cfg.max_gain_db / 20.0f);

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL,
        "audio_compressor_configure: rate=%d ch=%d block=%dms lookback=%dms lookahead=%dms maxGain=%.0fdB riseRate=%.1f/sec target=%.0fdBFS",
        ac->cfg.sample_rate, ac->cfg.channels, ac->cfg.block_duration_ms,
        ac->cfg.lookback_ms, ac->cfg.lookahead_ms, ac->cfg.max_gain_db,
        ac->cfg.rise_rate_per_sec, 20.0f * log10f(ac->cfg.target_level));
}

void audio_compressor_reset(struct audio_compressor* ac) {
    if (!ac) return;
    for (size_t i = 0; i < ac->pending_count; i++)
        u_free(ac->pending[i].samples);
    ac->pending_count = 0;
    ac->block_levels_count = 0;
    ac->accum_count = 0;
    ac->output_size = 0;
    ac->gain_smoothed = 1.0f;
    ac->dbg_counter = 0;
}

void audio_compressor_set_enabled(struct audio_compressor* ac, int enabled) {
    if (!ac) return;
    ac->enabled = enabled ? 1 : 0;
}

int audio_compressor_is_enabled(const struct audio_compressor* ac) {
    return ac ? ac->enabled : 0;
}

static float compute_block_level(const int16_t* samples, size_t count) {
    double sum_sq = 0.0;
    float peak = 0.0f;
    for (size_t i = 0; i < count; i++) {
        float v = (float)samples[i] / 32768.0f;
        sum_sq += (double)v * v;
        float av = fabsf(v);
        if (av > peak) peak = av;
    }
    float rms = (float)sqrt(sum_sq / (double)count);
    return (rms * 2.0f + peak) / 2.0f;
}

static void ensure_output_cap(struct audio_compressor* ac, size_t needed) {
    size_t want = ac->output_size + needed;
    if (want <= ac->output_cap) return;
    size_t new_cap = ac->output_cap ? ac->output_cap * 2 : 4096;
    while (new_cap < want) new_cap *= 2;
    int16_t* tmp = u_realloc(ac->output, new_cap * sizeof(int16_t));
    if (!tmp) return;
    ac->output = tmp;
    ac->output_cap = new_cap;
}

static void append_output(struct audio_compressor* ac, const int16_t* samples, size_t count) {
    ensure_output_cap(ac, count);
    memcpy(ac->output + ac->output_size, samples, count * sizeof(int16_t));
    ac->output_size += count;
}

static void ensure_block_levels_cap(struct audio_compressor* ac) {
    if (ac->block_levels_count < ac->block_levels_cap) return;
    size_t new_cap = ac->block_levels_cap ? ac->block_levels_cap * 2 : BLOCK_LEVELS_CHUNK;
    float* tmp = u_realloc(ac->block_levels, new_cap * sizeof(float));
    if (!tmp) return;
    ac->block_levels = tmp;
    ac->block_levels_cap = new_cap;
}

static void add_block_level(struct audio_compressor* ac, float level) {
    ensure_block_levels_cap(ac);
    ac->block_levels[ac->block_levels_count++] = level;
}

static void ensure_pending_cap(struct audio_compressor* ac) {
    if (ac->pending_count < ac->pending_cap) return;
    size_t new_cap = ac->pending_cap ? ac->pending_cap * 2 : PENDING_CHUNK;
    struct pending_block* tmp = u_realloc(ac->pending, new_cap * sizeof(*tmp));
    if (!tmp) return;
    ac->pending = tmp;
    ac->pending_cap = new_cap;
}

static void process_pending_block(struct audio_compressor* ac);

void audio_compressor_push(struct audio_compressor* ac, const int16_t* samples, size_t count) {
    if (!ac || !samples || count == 0) return;

    if (!ac->enabled) {
        append_output(ac, samples, count);
        return;
    }

    size_t remaining = count;
    const int16_t* src = samples;

    while (remaining > 0) {
        size_t need = (size_t)ac->block_samples - ac->accum_count;
        size_t take = remaining < need ? remaining : need;

        size_t new_cnt = ac->accum_count + take;
        int16_t* tmp = u_realloc(ac->accumulator, new_cnt * sizeof(int16_t));
        if (!tmp) return;
        ac->accumulator = tmp;
        memcpy(ac->accumulator + ac->accum_count, src, take * sizeof(int16_t));
        ac->accum_count = new_cnt;

        src += take;
        remaining -= take;

        if (ac->accum_count == (size_t)ac->block_samples) {
            float level = compute_block_level(ac->accumulator, ac->accum_count);
            add_block_level(ac, level);

            ensure_pending_cap(ac);
            ac->pending[ac->pending_count].samples    = ac->accumulator;
            ac->pending[ac->pending_count].count       = ac->accum_count;
            ac->pending[ac->pending_count].level_index = ac->block_levels_count - 1;
            ac->pending_count++;

            ac->accumulator = NULL;
            ac->accum_count = 0;
        }
    }

    while (ac->pending_count > (size_t)ac->lookahead_blocks)
        process_pending_block(ac);
}

static void process_pending_block(struct audio_compressor* ac) {
    if (ac->pending_count == 0) return;

    size_t idx = ac->pending[0].level_index;

    size_t win_start = idx >= (size_t)ac->lookback_blocks ? idx - (size_t)ac->lookback_blocks : 0;
    size_t win_end = idx + (size_t)ac->lookahead_blocks;
    if (win_end >= ac->block_levels_count) win_end = ac->block_levels_count - 1;

    float envelope = 0.0f;
    for (size_t i = win_start; i <= win_end; i++) {
        if (ac->block_levels[i] > envelope) envelope = ac->block_levels[i];
    }
    if (envelope < 0.0001f) envelope = 0.0001f;

    float G_raw = ac->cfg.target_level / envelope;
    float G_prev = ac->gain_smoothed;

    if (G_raw > ac->gain_smoothed) {
        float max_rise = ac->gain_smoothed * ac->rise_factor_per_block;
        ac->gain_smoothed = G_raw < max_rise ? G_raw : max_rise;
    } else {
        ac->gain_smoothed = G_raw;
    }
    if (ac->gain_smoothed > ac->max_gain) ac->gain_smoothed = ac->max_gain;

    float gain_db = 20.0f * log10f(ac->gain_smoothed > 0.0001f ? ac->gain_smoothed : 0.0001f);
    float prev_db = 20.0f * log10f(G_prev > 0.0001f ? G_prev : 0.0001f);
    if (ac->dbg_counter % 25 == 0 || fabsf(gain_db - prev_db) > 3.0f) {
        DEBUG_DEBUG(DEBUG_CATEGORY_GENERAL, "audio_compressor: block %zu level=%.4f envelope=%.4f gainRaw=%.1fdB gain=%.1fdB",
                    idx, ac->block_levels[idx], envelope, 20.0f * log10f(G_raw), gain_db);
    }
    ac->dbg_counter++;

    for (size_t i = 0; i < ac->pending[0].count; i++) {
        float v = (float)ac->pending[0].samples[i] * ac->gain_smoothed;
        if (v > 32767.0f) v = 32767.0f;
        if (v < -32768.0f) v = -32768.0f;
        ac->pending[0].samples[i] = (int16_t)(int)v;
    }
    append_output(ac, ac->pending[0].samples, ac->pending[0].count);
    u_free(ac->pending[0].samples);

    /* shift remaining pending blocks down */
    if (ac->pending_count > 1)
        memmove(ac->pending, ac->pending + 1, (ac->pending_count - 1) * sizeof(*ac->pending));
    ac->pending_count--;
}

void audio_compressor_flush(struct audio_compressor* ac) {
    if (!ac) return;
    if (!ac->enabled) return;

    if (ac->accum_count > 0) {
        float level = compute_block_level(ac->accumulator, ac->accum_count);
        add_block_level(ac, level);

        ensure_pending_cap(ac);
        ac->pending[ac->pending_count].samples    = ac->accumulator;
        ac->pending[ac->pending_count].count       = ac->accum_count;
        ac->pending[ac->pending_count].level_index = ac->block_levels_count - 1;
        ac->pending_count++;

        ac->accumulator = NULL;
        ac->accum_count = 0;
    }

    if (ac->pending_count > 0)
        DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "audio_compressor_flush: %zu pending blocks (gainSmoothed=%.1fdB)",
                   ac->pending_count, 20.0f * log10f(ac->gain_smoothed > 0.0001f ? ac->gain_smoothed : 0.0001f));

    while (ac->pending_count > 0)
        process_pending_block(ac);
}

const int16_t* audio_compressor_output(const struct audio_compressor* ac) {
    return ac ? ac->output : NULL;
}

size_t audio_compressor_output_size(const struct audio_compressor* ac) {
    return ac ? ac->output_size : 0;
}