#include <XCEngine/Audio/Equalizer.h>
#include <algorithm>
#include <cmath>

namespace XCEngine {
namespace Audio {

Equalizer::Equalizer()
    : m_sampleRate(48000)
{
    SetBandCount(4);

    m_frequencies[0] = 100.0f;
    m_frequencies[1] = 500.0f;
    m_frequencies[2] = 2000.0f;
    m_frequencies[3] = 8000.0f;

    for (uint32 i = 0; i < m_bandCount; ++i) {
        m_gains[i] = 0.0f;
        m_qs[i] = 1.0f;
        ComputeCoefficients(i, m_frequencies[i], m_qs[i], m_gains[i]);
    }
}

Equalizer::~Equalizer() {
}

void Equalizer::ProcessAudio(float* buffer, uint32 sampleCount, uint32 channels) {
    if (!m_enabled || buffer == nullptr || sampleCount == 0) {
        return;
    }

    if (channels == 0) {
        return;
    }

    for (uint32 ch = 0; ch < channels; ++ch) {
        for (uint32 band = 0; band < m_bandCount; ++band) {
            uint32 stateIndex = ch * m_bandCount + band;
            BandState& state = m_bandStates[stateIndex];

            float a0 = m_a0[band];
            float a1 = m_a1[band];
            float a2 = m_a2[band];
            float b1 = m_b1[band];
            float b2 = m_b2[band];

            for (uint32 i = 0; i < sampleCount; ++i) {
                uint32 sampleIndex = i * channels + ch;
                float x0 = buffer[sampleIndex];

                float y0 = a0 * x0 + a1 * state.x1 + a2 * state.x2 - b1 * state.y1 - b2 * state.y2;

                state.x2 = state.x1;
                state.x1 = x0;
                state.y2 = state.y1;
                state.y1 = y0;

                buffer[sampleIndex] = y0;
            }
        }
    }
}

void Equalizer::SetBandCount(uint32 count) {
    if (count == m_bandCount) {
        return;
    }

    m_bandCount = count;
    m_frequencies.resize(count, 1000.0f);
    m_gains.resize(count, 0.0f);
    m_qs.resize(count, 1.0f);
    m_a0.resize(count, 1.0f);
    m_a1.resize(count, 0.0f);
    m_a2.resize(count, 0.0f);
    m_b1.resize(count, 0.0f);
    m_b2.resize(count, 0.0f);

    m_bandStates.resize(count * 2);
}

void Equalizer::SetBandFrequency(uint32 band, float frequency) {
    if (band >= m_bandCount) {
        return;
    }

    m_frequencies[band] = std::max(20.0f, std::min(frequency, 20000.0f));
    ComputeCoefficients(band, m_frequencies[band], m_qs[band], m_gains[band]);
}

float Equalizer::GetBandFrequency(uint32 band) const {
    if (band >= m_bandCount) {
        return 0.0f;
    }
    return m_frequencies[band];
}

void Equalizer::SetBandGain(uint32 band, float gainDb) {
    if (band >= m_bandCount) {
        return;
    }

    m_gains[band] = std::max(-24.0f, std::min(gainDb, 24.0f));
    ComputeCoefficients(band, m_frequencies[band], m_qs[band], m_gains[band]);
}

float Equalizer::GetBandGain(uint32 band) const {
    if (band >= m_bandCount) {
        return 0.0f;
    }
    return m_gains[band];
}

void Equalizer::SetBandQ(uint32 band, float q) {
    if (band >= m_bandCount) {
        return;
    }

    m_qs[band] = std::max(0.1f, std::min(q, 10.0f));
    ComputeCoefficients(band, m_frequencies[band], m_qs[band], m_gains[band]);
}

float Equalizer::GetBandQ(uint32 band) const {
    if (band >= m_bandCount) {
        return 0.0f;
    }
    return m_qs[band];
}

void Equalizer::SetEnabled(bool enabled) {
    m_enabled = enabled;
}

void Equalizer::SetWetMix(float wetMix) {
    m_wetMix = std::max(0.0f, std::min(1.0f, wetMix));
}

void Equalizer::ComputeCoefficients(uint32 band, float frequency, float q, float gainDb) {
    if (band >= m_bandCount) {
        return;
    }

    float A = std::pow(10.0f, gainDb / 40.0f);
    float w0 = 2.0f * 3.14159265f * frequency / m_sampleRate;
    float cosw0 = std::cos(w0);
    float sinw0 = std::sin(w0);
    float alpha = sinw0 / (2.0f * q);

    float b0 = 1.0f + alpha * A;
    float b1 = -2.0f * cosw0;
    float b2 = 1.0f - alpha * A;
    float a0 = 1.0f + alpha / A;
    float a1 = -2.0f * cosw0;
    float a2 = 1.0f - alpha / A;

    m_a0[band] = b0 / a0;
    m_a1[band] = b1 / a0;
    m_a2[band] = b2 / a0;
    m_b1[band] = a1 / a0;
    m_b2[band] = a2 / a0;
}

} // namespace Audio
} // namespace XCEngine