// SPDX-FileCopyrightText: 2024 Gary Wang <git@blumia.net>
//
// SPDX-License-Identifier: MIT

#include "fftspectrum.h"

#include <QAudioBuffer>
#include <QAudioBufferOutput>
#include <QPainter>

#include <kissfft.hh>

FFTSpectrum::FFTSpectrum(QWidget* parent)
    : QWidget(parent)
    , m_audioBufferOutput(new QAudioBufferOutput(this))
{
    connect(this, &FFTSpectrum::mediaPlayerChanged, this, [=]() {
        if (m_mediaPlayer) {
            m_mediaPlayer->setAudioBufferOutput(m_audioBufferOutput);
        }
    });

    connect(m_audioBufferOutput, &QAudioBufferOutput::audioBufferReceived, this, [=](const QAudioBuffer& buffer) {
        const QAudioFormat&& fmt = buffer.format();
        const QAudioFormat::SampleFormat sampleFormat = fmt.sampleFormat();
        int channelCount = fmt.channelCount();
        const int frameCount = buffer.frameCount();
        kissfft<float> fft(frameCount, false);
        std::vector<kissfft<float>::cpx_t> samples(frameCount);
        std::vector<kissfft<float>::cpx_t> mass(frameCount);
        if (sampleFormat == QAudioFormat::Int16 && channelCount == 1) {
            const QAudioBuffer::S16M* data = buffer.constData<QAudioBuffer::S16M>();
            for (int i = 0; i < frameCount; ++i) {
                samples[i].real(data[i].value(QAudioFormat::FrontCenter) / float(32768));
                samples[i].imag(0);
            }
        } else if (sampleFormat == QAudioFormat::Float && channelCount == 2) {
            const QAudioBuffer::F32S* data = buffer.constData<QAudioBuffer::F32S>();
            for (int i = 0; i < frameCount; ++i) {
                samples[i].real(data[i].value(QAudioFormat::FrontLeft));
                samples[i].imag(0);
            }
        }
        fft.transform(samples.data(), mass.data());
        m_freq.resize(frameCount);
        for (int i = 0; i < frameCount; i++) {
            m_freq[i] = sqrt(mass[i].real() * mass[i].real() + mass[i].imag() * mass[i].imag());
        }
        for (auto& val : m_freq) {
            val = log10(val + 1);
        }
        qDebug() << m_freq.size() << "updated";
        update();
        qDebug() << sampleFormat << channelCount;
    });

    resize(490, 420);
}

FFTSpectrum::~FFTSpectrum()
{

}

void FFTSpectrum::setMediaPlayer(QMediaPlayer* player)
{
    m_mediaPlayer = player;
    emit mediaPlayerChanged();
}

void FFTSpectrum::paintEvent(QPaintEvent* e)
{
    QPainter painter(this);
    painter.fillRect(rect(), Qt::black);

    if (!m_freq.empty()) {
        int width = this->width();
        int height = this->height();
        int barWidth = std::max(1ULL, width / m_freq.size());
        qDebug() << m_freq.size() << barWidth << m_freq << "freq";

        for (int i = 0; i < m_freq.size(); i++) {
            int barHeight = static_cast<int>(m_freq[i] * height);
            painter.fillRect(i * barWidth, height - barHeight, barWidth, barHeight, Qt::green);
        }
    }
}