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

/* История и задержанные блоки ограничены окном настройки. Только запись накапливает output. */
struct audio_compressor {
    audio_compressor_config_t cfg;
    int enabled, configured, failed, block_samples, lookback_blocks, lookahead_blocks;
    float rise_factor_per_block, release_factor_per_block, max_gain, gain_smoothed;
    float* block_levels;
    size_t levels_cap, blocks_seen;
    int16_t* pending;
    size_t pending_cap, pending_head, pending_count, accum_count;
    int16_t* output;
    size_t output_size, output_cap;
    uint32_t clip_count;
    unsigned dbg_counter;
};

/* Начальное усиление в середине диапазона в дБ. */
static float initial_gain(float max_gain) { return max_gain > 0.0f ? sqrtf(max_gain) : 1.0f; }

struct audio_compressor* audio_compressor_create(void) {
    struct audio_compressor* ac = u_calloc(1, sizeof(*ac));
    if (!ac) { DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: state allocation failed"); return NULL; }
    ac->enabled = 1;
    return ac;
}

void audio_compressor_destroy(struct audio_compressor* ac) {
    if (!ac) return;
    u_free(ac->pending); u_free(ac->block_levels); u_free(ac->output); u_free(ac);
}

void audio_compressor_reset(struct audio_compressor* ac) {
    if (!ac) return;
    ac->pending_head = ac->pending_count = ac->accum_count = ac->blocks_seen = ac->output_size = 0;
    ac->gain_smoothed = initial_gain(ac->max_gain);
    ac->clip_count = ac->dbg_counter = 0;
    ac->failed = 0;
}

/* Ноль lookahead означает отсутствие задержки; defaults остальных полей задаются явно здесь. */
int audio_compressor_configure(struct audio_compressor* ac, const audio_compressor_config_t* config) {
    if (!ac || !config) { DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: missing configuration"); return -1; }
    audio_compressor_config_t cfg = *config;
    if (!cfg.sample_rate) cfg.sample_rate = 48000;
    if (!cfg.channels) cfg.channels = 1;
    if (!cfg.block_duration_ms) cfg.block_duration_ms = 20;
    if (!cfg.lookback_ms) cfg.lookback_ms = 200;
    if (!cfg.target_level) cfg.target_level = 0.25f;
    if (!cfg.rise_rate_per_sec) cfg.rise_rate_per_sec = 4.0f;
    if (!cfg.release_rate_per_sec) cfg.release_rate_per_sec = cfg.rise_rate_per_sec * 5.0f;
    if (cfg.sample_rate < 8000 || cfg.sample_rate > 192000 || cfg.channels < 1 || cfg.channels > 2 ||
        cfg.block_duration_ms < 1 || cfg.block_duration_ms > 1000 || cfg.lookback_ms < 0 || cfg.lookback_ms > 5000 ||
        cfg.lookahead_ms < 0 || cfg.lookahead_ms > 5000 || !isfinite(cfg.target_level) || cfg.target_level <= 0 ||
        !isfinite(cfg.max_gain_db) || cfg.max_gain_db < 0 || cfg.max_gain_db > 60 ||
        !isfinite(cfg.rise_rate_per_sec) || cfg.rise_rate_per_sec <= 0 ||
        !isfinite(cfg.release_rate_per_sec) || cfg.release_rate_per_sec <= 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: invalid rate=%d ch=%d block=%d back=%d ahead=%d",
                    cfg.sample_rate, cfg.channels, cfg.block_duration_ms, cfg.lookback_ms, cfg.lookahead_ms);
        return -1;
    }
    int block = cfg.sample_rate * cfg.block_duration_ms / 1000 * cfg.channels;
    int back = cfg.lookback_ms / cfg.block_duration_ms, ahead = cfg.lookahead_ms / cfg.block_duration_ms;
    size_t levels_cap = (size_t)back + ahead + 2, pending_cap = (size_t)ahead + 2;
    float* levels = u_calloc(levels_cap, sizeof(float));
    int16_t* pending = u_calloc(pending_cap * (size_t)block, sizeof(int16_t));
    if (!levels || !pending) {
        u_free(levels); u_free(pending);
        DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: buffers allocation failed");
        return -1;
    }
    u_free(ac->block_levels); u_free(ac->pending);
    ac->cfg = cfg; ac->block_samples = block;
    ac->lookback_blocks = back; ac->lookahead_blocks = ahead;
    ac->block_levels = levels; ac->levels_cap = levels_cap;
    ac->pending = pending; ac->pending_cap = pending_cap;
    ac->rise_factor_per_block = powf(cfg.rise_rate_per_sec, (float)cfg.block_duration_ms / 1000.0f);
    ac->release_factor_per_block = powf(cfg.release_rate_per_sec, (float)cfg.block_duration_ms / 1000.0f);
    ac->max_gain = powf(10.0f, cfg.max_gain_db / 20.0f);
    ac->configured = 1;
    audio_compressor_reset(ac);
    DEBUG_INFO(DEBUG_CATEGORY_MEDIA, "compressor: configured rate=%d ch=%d ahead=%dms history=%zu pending=%zu",
               cfg.sample_rate, cfg.channels, cfg.lookahead_ms, levels_cap, pending_cap);
    return 0;
}

struct audio_compressor* audio_compressor_clone_config(const struct audio_compressor* source) {
    if (!source || !source->configured) { DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: cannot clone configuration"); return NULL; }
    struct audio_compressor* copy = audio_compressor_create();
    if (copy && audio_compressor_configure(copy, &source->cfg) != 0) { audio_compressor_destroy(copy); return NULL; }
    if (copy) copy->enabled = source->enabled;
    return copy;
}

void audio_compressor_set_enabled(struct audio_compressor* ac, int enabled) { if (ac) ac->enabled = !!enabled; }
int audio_compressor_is_enabled(const struct audio_compressor* ac) { return ac && ac->enabled; }

/* Одна оценка для всех каналов сохраняет баланс стерео. */
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;
        if (fabsf(v) > peak) peak = fabsf(v);
    }
    return (2.0f * sqrtf((float)(sum_sq / count)) + peak) / 2.0f;
}

/* Выход записи растёт только вне аппаратного callback; OOM переводит запись в failed. */
static int reserve_output(struct audio_compressor* ac, size_t count) {
    if (count > SIZE_MAX / sizeof(int16_t) - ac->output_size) goto failed;
    size_t want = ac->output_size + count;
    if (want <= ac->output_cap) return 0;
    size_t capacity = ac->output_cap ? ac->output_cap : 4096;
    while (capacity < want) {
        if (capacity > SIZE_MAX / sizeof(int16_t) / 2) { capacity = want; break; }
        capacity *= 2;
    }
    int16_t* samples = u_realloc(ac->output, capacity * sizeof(int16_t));
    if (!samples) goto failed;
    ac->output = samples; ac->output_cap = capacity;
    return 0;
failed:
    ac->failed = 1;
    DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: recording output allocation failed samples=%zu", count);
    return -1;
}

/* Общий алгоритм усиления для записи с lookahead и потокового кадра без lookahead. */
static void apply_gain(struct audio_compressor* ac, int16_t* samples, size_t count, size_t idx) {
    size_t begin = idx >= (size_t)ac->lookback_blocks ? idx - ac->lookback_blocks : 0;
    size_t end = idx + ac->lookahead_blocks;
    if (end >= ac->blocks_seen) end = ac->blocks_seen - 1;
    float envelope = 0;
    for (size_t i = begin; i <= end; i++) {
        float level = ac->block_levels[i % ac->levels_cap];
        if (level > envelope) envelope = level;
    }
    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) {
        /* сигнал стал тише — плавно поднимаем gain (атака) */
        float max_rise = ac->gain_smoothed * ac->rise_factor_per_block;
        ac->gain_smoothed = G_raw < max_rise ? G_raw : max_rise;
    } else {
        /* сигнал стал громче — плавно опускаем gain (release), без мгновенной ступеньки
         * (мгновенный дроп давал клик на границе 20мс-блока). */
        float min_fall = ac->gain_smoothed / ac->release_factor_per_block;
        ac->gain_smoothed = G_raw > min_fall ? G_raw : min_fall;
    }
    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 clip=%u",
                    idx, ac->block_levels[idx % ac->levels_cap], envelope, 20.0f * log10f(G_raw), gain_db, (unsigned)ac->clip_count);
    }
    ac->dbg_counter++;

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

/* Слить один готовый блок записи. История хранится по абсолютному номеру в кольце. */
static int drain_block(struct audio_compressor* ac, size_t count) {
    if (reserve_output(ac, count) != 0) return -1;
    int16_t* samples = ac->pending + ac->pending_head * ac->block_samples;
    apply_gain(ac, samples, count, ac->blocks_seen - ac->pending_count);
    memcpy(ac->output + ac->output_size, samples, count * sizeof(int16_t));
    ac->output_size += count;
    ac->pending_head = (ac->pending_head + 1) % ac->pending_cap;
    --ac->pending_count;
    return 0;
}

/* Ровно один кадр, без output-накопителя и выделений памяти. in/out могут совпадать. */
int audio_compressor_process_frame(struct audio_compressor* ac, const int16_t* in, int16_t* out, size_t count) {
    if (!ac || !ac->configured || ac->failed || !in || !out || count != (size_t)ac->block_samples ||
        ac->lookahead_blocks || ac->pending_count || ac->accum_count) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: invalid streaming frame count=%zu", count);
        return -1;
    }
    memmove(out, in, count * sizeof(int16_t));
    if (ac->enabled) {
        size_t idx = ac->blocks_seen++;
        ac->block_levels[idx % ac->levels_cap] = compute_block_level(in, count);
        apply_gain(ac, out, count, idx);
    }
    return 0;
}

int audio_compressor_push(struct audio_compressor* ac, const int16_t* samples, size_t count) {
    if (!ac || !ac->configured || ac->failed || (!samples && count)) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: invalid recording input"); return -1;
    }
    if (!ac->enabled) {
        if (reserve_output(ac, count) != 0) return -1;
        if (count) memcpy(ac->output + ac->output_size, samples, count * sizeof(int16_t));
        ac->output_size += count;
        return 0;
    }
    while (count) {
        size_t take = (size_t)ac->block_samples - ac->accum_count;
        if (take > count) take = count;
        int16_t* block = ac->pending + ((ac->pending_head + ac->pending_count) % ac->pending_cap) * ac->block_samples;
        memcpy(block + ac->accum_count, samples, take * sizeof(int16_t));
        samples += take; count -= take; ac->accum_count += take;
        if (ac->accum_count == (size_t)ac->block_samples) {
            ac->block_levels[ac->blocks_seen++ % ac->levels_cap] = compute_block_level(block, ac->accum_count);
            ac->accum_count = 0; ++ac->pending_count;
            if (ac->pending_count > (size_t)ac->lookahead_blocks && drain_block(ac, ac->block_samples) != 0) return -1;
        }
    }
    return 0;
}

int audio_compressor_flush(struct audio_compressor* ac) {
    if (!ac || !ac->configured || ac->failed) { DEBUG_ERROR(DEBUG_CATEGORY_MEDIA, "compressor: invalid flush"); return -1; }
    if (!ac->enabled) return 0;
    size_t tail = ac->accum_count;
    if (tail) {
        int16_t* block = ac->pending + ((ac->pending_head + ac->pending_count) % ac->pending_cap) * ac->block_samples;
        ac->block_levels[ac->blocks_seen++ % ac->levels_cap] = compute_block_level(block, tail);
        ac->accum_count = 0; ++ac->pending_count;
    }
    while (ac->pending_count) {
        size_t count = tail && ac->pending_count == 1 ? tail : (size_t)ac->block_samples;
        if (drain_block(ac, count) != 0) return -1;
    }
    return 0;
}

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; }
void audio_compressor_clear_output(struct audio_compressor* ac) { if (ac) ac->output_size = 0; }
uint32_t audio_compressor_clip_count(const struct audio_compressor* ac) { return ac ? ac->clip_count : 0; }