#include "spectral.hpp"
#include "fftconv.hpp"
#include <cmath>
#include <cstring>
#include <vector>
#include <cstdlib>

SpectralProcessor::SpectralProcessor(size_t nfft, size_t hop, float sample_rate)
    : nfft_(nfft), hop_(hop), frame_count_(0), output_pos_(0),
      detector_(nfft, sample_rate) {
    window_ = new double[nfft_];
    computeWindow();
    fft::init_plan(&plan_, static_cast<uint32_t>(std::log2(nfft_)));
    buf_ = new std::complex<double>[nfft_];
    tmp_buf_ = new std::complex<double>[nfft_];
    fir_buf_ = new std::complex<double>[nfft_];
    fir_freq_ = new std::complex<double>[nfft_];
    overlap_.resize(nfft_, 0.0f);
    mask_.resize(nfft_, 1.0f);

    // Build FIR window: falling half of periodic Hann(4096).
    // Plugin reads window[N/2..N-1] of periodic Hann (rising 0→1).
    fir_window_.resize(nfft_);
    for (size_t i = 0; i < nfft_; i++) {
        fir_window_[i] = 0.5 * (1.0 - std::cos(2.0 * M_PI * i / nfft_));
    }
}

SpectralProcessor::~SpectralProcessor() {
    delete[] window_;
    delete[] buf_;
    delete[] tmp_buf_;
    delete[] fir_buf_;
    delete[] fir_freq_;
}

void SpectralProcessor::setDetectorParams(const std::vector<DetectorBand>& bands) {
    detector_.setParams(bands);
}

void SpectralProcessor::computeWindow() {
    for (size_t i = 0; i < nfft_; i++) {
        window_[i] = 0.5 * (1.0 - std::cos(2.0 * M_PI * i / (nfft_ - 1)));
    }
}

void SpectralProcessor::stftFrame(const float* in, std::complex<double>* out) {
    for (size_t i = 0; i < nfft_; i++) {
        out[i] = std::complex<double>(static_cast<double>(in[i]) * window_[i], 0.0);
    }
    fft::execute(&plan_, out);
}

void SpectralProcessor::istftFrame(std::complex<double>* in, float* out, float* overlap) {
    memcpy(tmp_buf_, in, nfft_ * sizeof(std::complex<double>));
    fft::execute_inverse(&plan_, tmp_buf_);
    static bool wola_computed = false;
    static float wola_norm = 1.0f;
    if (!wola_computed) {
        double wola_sum = 0.0;
        for (size_t i = 0; i < nfft_; i++) {
            wola_sum += window_[i] * window_[i];
        }
        wola_norm = static_cast<float>(wola_sum / hop_);
        wola_computed = true;
    }
    for (size_t i = 0; i < nfft_; i++) {
        overlap[i] += static_cast<float>(tmp_buf_[i].real() * window_[i]);
    }
    for (size_t i = 0; i < hop_; i++) {
        out[i] = overlap[i] / wola_norm;
    }
    for (size_t i = 0; i < nfft_ - hop_; i++) {
        overlap[i] = overlap[i + hop_];
    }
    for (size_t i = nfft_ - hop_; i < nfft_; i++) {
        overlap[i] = 0.0f;
    }
}

void SpectralProcessor::processBlock(float* in, float* out, size_t num_samples, size_t num_channels) {
    memset(out, 0, num_samples * sizeof(float));
    if (num_samples == 0 || num_samples < nfft_) {
        return;
    }

    static const int firconv = []() {
        const char* e = getenv("RT_FIRCONV");
        return e ? atoi(e) : 0;
    }();

    size_t nframes = (num_samples - nfft_) / hop_ + 1;

    for (size_t f = 0; f < nframes; f++) {
        size_t offset = f * hop_;
        if (offset + nfft_ > num_samples) break;
        stftFrame(in + offset, buf_);

        detector_.processFrame(buf_, mask_.data());

        if (firconv) {
            // RT_FIRCONV=1: Build FIR from mask and apply via complex multiply.
            // The mask is real-valued (per-bin gain). We apply it directly
            // to the audio spectrum via complex multiply (th_b3c0 equivalent).
            // No upper-half zeroing — preserve Hermitian symmetry.
            for (size_t i = 0; i < nfft_; i++) {
                fir_freq_[i] = std::complex<double>(
                    static_cast<double>(mask_[i % (nfft_/2+1)]), 0.0);
            }

            // Complex multiply FIR × audio spectrum.
            for (size_t i = 0; i < nfft_; i++) {
                buf_[i] *= fir_freq_[i];
            }
        } else {
            // Default path: simple frequency-domain mask multiply.
            for (size_t i = 0; i < nfft_; i++) {
                buf_[i] *= mask_[i];
            }
        }

        istftFrame(buf_, out + offset, overlap_.data());
    }
}