#include "SpectrogramWidget.h"
#include "GradientDialog.h"
#include <QTimer>
#include <QSurfaceFormat>
#include <QFile>
#include <QMouseEvent>
#include <QJsonDocument>
#include <QJsonArray>
#include <QJsonObject>
#include <QDir>
#include <QStandardPaths>
#include <cmath>

SpectrogramWidget::SpectrogramWidget(QWidget *parent)
    : QOpenGLWidget(parent)
{
    setAttribute(Qt::WA_OpaquePaintEvent);

    QSurfaceFormat fmt;
    fmt.setVersion(3, 3);
    fmt.setProfile(QSurfaceFormat::CoreProfile);
    fmt.setSwapInterval(1);
    fmt.setSwapBehavior(QSurfaceFormat::DoubleBuffer);
    setFormat(fmt);

    setMinimumSize(400, 200);
    setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Expanding);

    loadBarColors();

    auto *timer = new QTimer(this);
    connect(timer, &QTimer::timeout,
            this, QOverload<>::of(&SpectrogramWidget::update));
    timer->start(16);

    m_frameTimer.start();
}

SpectrogramWidget::~SpectrogramWidget()
{
    makeCurrent();
    glDeleteTextures(1, &m_heatmapTex);
    glDeleteVertexArrays(1, &m_vaoBars);
    doneCurrent();
}

void SpectrogramWidget::receiveFft(const QVector<float> &mag)
{
    QMutexLocker lk(&m_latestMtx);
    m_latestMag          = mag;
    m_newMagAvailable    = true;
}

void SpectrogramWidget::reset()
{
    {
        QMutexLocker lk(&m_latestMtx);
        m_latestMag.clear();
        m_newMagAvailable = false;
    }
    m_currentMag.fill(0.f);
    m_barSmoothed.fill(0.f);
    m_barPeak.fill(0.f);
    m_peakHoldTimer.fill(0.f);
    m_dirty = false;
    update();
}

void SpectrogramWidget::initializeGL()
{
    initializeOpenGLFunctions();
    initShaders();
    updateHeatmapTexture();

    glGenVertexArrays(1, &m_vaoBars);
    glDisable(GL_BLEND);

    const int    bins   = engine::SpectrumAnalyzer::BIN_COUNT;
    const double logMin = std::log10(20.0);
    const double logMax = std::log10(22050.0);
    const double range  = logMax - logMin;

    for (int b = 0; b < BAR_COUNT; ++b) {
        double f1 = std::pow(10.0, logMin + range *  b      / BAR_COUNT);
        double f2 = std::pow(10.0, logMin + range * (b + 1) / BAR_COUNT);
        m_binMap[b].lo = std::clamp(int(f1 / 22050.0 * bins), 0, bins - 1);
        m_binMap[b].hi = std::clamp(int(f2 / 22050.0 * bins), m_binMap[b].lo, bins - 1);
    }
}

void SpectrogramWidget::initShaders()
{
    auto load = [](const QString &path) -> QByteArray {
        QFile f(path);
        if (!f.open(QIODevice::ReadOnly | QIODevice::Text)) {
            qWarning() << "Cannot open shader:" << path;
            return {};
        }
        return f.readAll();
    };

    if (!m_barsProg.addShaderFromSourceCode(QOpenGLShader::Vertex,
                                            load(":/shaders/bars.vert")))
        qWarning() << "Bars vert error:" << m_barsProg.log();
    if (!m_barsProg.addShaderFromSourceCode(QOpenGLShader::Fragment,
                                            load(":/shaders/bars.frag")))
        qWarning() << "Bars frag error:" << m_barsProg.log();
    if (!m_barsProg.link())
        qWarning() << "Bars link error:" << m_barsProg.log();

    m_uBarHeights = m_barsProg.uniformLocation("barHeights");
    m_uBarPeaks   = m_barsProg.uniformLocation("barPeaks");
    m_uBarCount   = m_barsProg.uniformLocation("barCount");
    m_uBarWidth   = m_barsProg.uniformLocation("barWidth");
    m_uGap        = m_barsProg.uniformLocation("gap");
    m_uBottomY    = m_barsProg.uniformLocation("bottomY");
    m_uTopClip    = m_barsProg.uniformLocation("topClip");
    m_uHeatmap    = m_barsProg.uniformLocation("heatmap");
}

void SpectrogramWidget::updateHeatmapTexture()
{
    if (m_stops.size() < 2)
        m_stops = { QColor(0, 0, 0), QColor(255, 255, 255) };

    const int nSegs = m_stops.size() - 1;
    auto lerp = [](int x, int y, float f) {
        return std::clamp(int(x + (y - x) * f), 0, 255);
    };
    for (int i = 0; i < 256; ++i) {
        float  t   = float(i) / 255.f * nSegs;
        int    seg = std::min(int(t), nSegs - 1);
        float  f   = t - seg;
        QColor a   = m_stops[seg], b = m_stops[seg + 1];
        m_heatmap[i] = qRgb(lerp(a.red(),   b.red(),   f),
                            lerp(a.green(), b.green(), f),
                            lerp(a.blue(),  b.blue(),  f));
    }

    std::array<GLubyte, 256 * 3> td;
    for (int i = 0; i < 256; ++i) {
        td[3 * i + 0] = static_cast<GLubyte>(qRed  (m_heatmap[i]));
        td[3 * i + 1] = static_cast<GLubyte>(qGreen(m_heatmap[i]));
        td[3 * i + 2] = static_cast<GLubyte>(qBlue (m_heatmap[i]));
    }

    if (m_heatmapTex) {
        glDeleteTextures(1, &m_heatmapTex);
        m_heatmapTex = 0;
    }
    glGenTextures(1, &m_heatmapTex);
    glBindTexture(GL_TEXTURE_1D, m_heatmapTex);
    glTexImage1D(GL_TEXTURE_1D, 0, GL_RGB, 256, 0,
                 GL_RGB, GL_UNSIGNED_BYTE, td.data());
    glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_WRAP_S,     GL_CLAMP_TO_EDGE);
}

void SpectrogramWidget::setStops(const QVector<QColor> &stops)
{
    if (stops.size() < 2) return;
    m_stops = stops;
    makeCurrent();
    updateHeatmapTexture();
    doneCurrent();
    update();
    saveBarColors();
}

void SpectrogramWidget::loadBarColors()
{
    QString path = QStandardPaths::writableLocation(QStandardPaths::HomeLocation)
                   + "/.subwave/barcolors.json";
    QFile f(path);
    if (!f.open(QIODevice::ReadOnly)) return;

    QJsonDocument doc = QJsonDocument::fromJson(f.readAll());
    if (!doc.isArray()) return;

    QVector<QColor> loaded;
    for (const auto &v : doc.array()) {
        QJsonObject obj = v.toObject();
        loaded.append(QColor(obj["r"].toInt(), obj["g"].toInt(), obj["b"].toInt()));
    }
    if (loaded.size() >= 2)
        m_stops = loaded;
}

void SpectrogramWidget::saveBarColors()
{
    QString dir  = QStandardPaths::writableLocation(QStandardPaths::HomeLocation)
                   + "/.subwave";
    QDir().mkpath(dir);
    QFile f(dir + "/barcolors.json");
    if (!f.open(QIODevice::WriteOnly)) return;

    QJsonArray arr;
    for (const QColor &c : m_stops) {
        QJsonObject obj;
        obj["r"] = c.red();
        obj["g"] = c.green();
        obj["b"] = c.blue();
        arr.append(obj);
    }
    f.write(QJsonDocument(arr).toJson());
}


void SpectrogramWidget::mousePressEvent(QMouseEvent *e)
{
    if (e->button() == Qt::LeftButton) {
        GradientDialog dlg(m_stops, this);
        if (dlg.exec() == QDialog::Accepted)
            setStops(dlg.stops());
    }
    QOpenGLWidget::mousePressEvent(e);
}


void SpectrogramWidget::computeBars(float dt)
{
    const int fftLen = m_currentMag.size();

    for (int b = 0; b < BAR_COUNT; ++b) {
        const int lo = m_binMap[b].lo;
        const int hi = std::min(m_binMap[b].hi, fftLen - 1);

        float sum = 0.f;
        int   cnt = 0;
        for (int k = lo; k <= hi; ++k) { sum += m_currentMag[k]; ++cnt; }
        float raw = cnt > 0 ? sum / cnt : 0.f;

        float db  = raw > 0.f ? 20.f * std::log10(raw) : DB_MIN;
        float bar = std::clamp((db - DB_MIN) / (DB_MAX - DB_MIN), 0.f, 1.f);

        const float coeff = bar > m_barSmoothed[b] ? 0.50f : 0.18f;
        m_barSmoothed[b] += (bar - m_barSmoothed[b]) * coeff;

        if (m_barSmoothed[b] >= m_barPeak[b]) {
            m_barPeak[b]       = m_barSmoothed[b];
            m_peakHoldTimer[b] = PEAK_HOLD_SEC;
        } else if (m_peakHoldTimer[b] > 0.f) {
            m_peakHoldTimer[b] -= dt;
        } else {
            m_barPeak[b] = std::max(0.f, m_barPeak[b] - PEAK_DECAY_SEC * dt);
        }

        m_barHeights [b] = m_barSmoothed[b];
        m_barPeaksVec[b] = m_barPeak[b];
    }
}


void SpectrogramWidget::paintGL()
{
    const float dt = static_cast<float>(m_frameTimer.restart()) / 1000.f;

    {
        QMutexLocker lk(&m_latestMtx);
        if (m_newMagAvailable) {
            m_currentMag      = std::move(m_latestMag);
            m_latestMag       = {};
            m_newMagAvailable = false;
            m_dirty           = true;
        }
    }

    computeBars(dt);

    QColor bg = palette().color(QPalette::Window);
    glClearColor(bg.redF(), bg.greenF(), bg.blueF(), 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);

    if (!m_barsProg.isLinked()) return;

    const float barWidth  = 2.0f / BAR_COUNT;
    const float gap       = barWidth * 0.2f;
    const float drawWidth = barWidth - gap;

    m_barsProg.bind();
    glUniform1fv(m_uBarHeights, BAR_COUNT, m_barHeights .data());
    glUniform1fv(m_uBarPeaks,   BAR_COUNT, m_barPeaksVec.data());
    glUniform1i (m_uBarCount,   BAR_COUNT);
    glUniform1f (m_uBarWidth,   drawWidth);
    glUniform1f (m_uGap,        gap);
    glUniform1f (m_uBottomY,   -1.0f);
    glUniform1f (m_uTopClip,    1.0f);
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_1D, m_heatmapTex);
    glUniform1i(m_uHeatmap, 0);
    glBindVertexArray(m_vaoBars);
    glDrawArraysInstanced(GL_TRIANGLES, 0, 6, BAR_COUNT);
    glBindVertexArray(0);
    m_barsProg.release();
}