// Copyright © 2025 Bjango. All rights reserved.

#pragma once

#include "LPCvocoder.hpp"
#include "AudioUtils.hpp"
#include "PeakMeter.hpp"
#import "ParameterAddresses.h"
#import <AudioToolbox/AudioToolbox.h>
#import <algorithm>
#import <vector>
#import <span>
#import <memory>
#include <cstring>

// Plain struct because std::atomic isn't compatible with Swift's C++ interop.
// Access is through __atomic_* builtins in DSPKernel's methods.
struct DSPSharedState {
    static constexpr uint32_t kParamCount = 15;
    uint32_t paramValues[kParamCount] = {};  // float bits stored as uint32_t for __atomic_* compatibility
    uint32_t dirtyMask = 0;

    // Meter values (float bits stored as uint32_t)
    uint32_t synthLevel = 0;
    uint32_t voiceLevel = 0;
    uint32_t outLevel = 0;
    uint32_t meterDirty = 0;
};

class DSPKernel {
public:
    uint32_t consumeDirtyParamMask() {
        return __atomic_exchange_n(&sharedState.dirtyMask, 0u, __ATOMIC_ACQ_REL);
    }

    float getAtomicParamValue(uint32_t address) const {
        if (address < DSPSharedState::kParamCount) {
            return loadFloat(sharedState.paramValues[address]);
        }
        return 0.0f;
    }

    bool consumeMeterDirty() {
        return __atomic_exchange_n(&sharedState.meterDirty, 0u, __ATOMIC_ACQ_REL) != 0;
    }

    float getAtomicMeterSynthLevel() const { return loadFloat(sharedState.synthLevel); }
    float getAtomicMeterVoiceLevel() const { return loadFloat(sharedState.voiceLevel); }
    float getAtomicMeterOutLevel()   const { return loadFloat(sharedState.outLevel); }
    
    void initialize(double inSampleRate, bool voiceOnlyOutput, float wetMix, float dryMix, float compressorRatio, float compressorThreshold, float compressorAttackMs, float compressorReleaseMs) {
        sampleRate = inSampleRate;
        
        processor = std::make_unique<LPCVocoderProcessor>(inSampleRate);
        processor->setQuality(quality);
        processor->setWhiteNoiseLevel(whiteNoiseLevelDb);
        processor->setWhiteNoiseThreshold(whiteNoiseThreshold);
        processor->setOutputGain(outputGainDb);
        processor->setVoiceInGain(voiceInGainDb);
        processor->setSynthInGain(synthInGainDb);
        processor->setVoiceMode(voiceMode);
        processor->setFormantShift(formantShiftCents);
        processor->setHighPassFilter(highPassFilter);
        processor->setLowPassFilter(lowPassFilter);
        processor->setVoiceFreeze(voiceFreeze);
        processor->setDownsampleTo31Khz(crush);
        processor->set16Bit(crush);
        processor->setCompressEnabled(compress);
        
        // hard coded flags
        processor->setVoiceOnlyOutput(voiceOnlyOutput);
        processor->setWetMix(wetMix);
        processor->setDryMix(dryMix);
        processor->setCompressorSettings(compressorRatio, compressorThreshold, compressorAttackMs, compressorReleaseMs);
        
        peakMeter = std::make_unique<PeakMeter>();
    }

    void deInitialize() {
        processor.reset();
    }

    bool isBypassed() {
        return bypassed;
    }

    void setBypass(bool shouldBypass) {
        bypassed = shouldBypass;
    }
    
    void setLicensed(bool isLicensed) {
        licensed = isLicensed;
    }

    // MARK: - Parameter Getter / Setter
    void setParameter(AUParameterAddress address, AUValue value, bool markDirty = false) {
        switch (address) {
            case RobotRocketParameterAddress::quality:
                quality = value;
                if (processor) {
                    processor->setQuality(value);
                }
                break;
            case RobotRocketParameterAddress::whiteNoiseLevel:
                whiteNoiseLevelDb = value;
                if (processor) {
                    processor->setWhiteNoiseLevel(value);
                }
                break;
            case RobotRocketParameterAddress::whiteNoiseThreshold:
                whiteNoiseThreshold = value;
                if (processor) {
                    processor->setWhiteNoiseThreshold(value);
                }
                
                break;
            case RobotRocketParameterAddress::outputLevel:
                outputGainDb = value;
                if (processor) {
                    processor->setOutputGain(value);
                }
                
                break;
            case RobotRocketParameterAddress::voiceInGain:
                voiceInGainDb = value;
                if (processor) {
                    processor->setVoiceInGain(value);
                }
                if (peakMeter) {
                    float voiceInGainFactor = AudioUtils::decibelToScale(value);
                    peakMeter->setVoiceGain(voiceInGainFactor);
                }
                
                break;
            case RobotRocketParameterAddress::synthInGain:
                synthInGainDb = value;
                if (processor) {
                    processor->setSynthInGain(value);
                }
                if (peakMeter) {
                    float synthInGainFactor = AudioUtils::decibelToScale(value);
                    peakMeter->setSynthGain(synthInGainFactor);
                }
                
                break;
            case RobotRocketParameterAddress::voiceMode:
                voiceMode = static_cast<int>(value);
                if (processor) {
                    processor->setVoiceMode(voiceMode);
                }
                
                break;
            case RobotRocketParameterAddress::formantShift:
                formantShiftCents = value;
                if (processor) {
                    processor->setFormantShift(value);
                }
                
                break;
            case RobotRocketParameterAddress::highPassFilter:
                highPassFilter = (value > 0.5f);
                if (processor) {
                    processor->setHighPassFilter(highPassFilter);
                }
                
                break;
            case RobotRocketParameterAddress::lowPassFilter:
                lowPassFilter = (value > 0.5f);
                if (processor) {
                    processor->setLowPassFilter(lowPassFilter);
                }
                
                break;
            case RobotRocketParameterAddress::freeze:
                voiceFreeze = (value > 0.5f);
                if (processor) {
                    processor->setVoiceFreeze(voiceFreeze);
                }
                
                break;
            case RobotRocketParameterAddress::crush:
                crush = (value > 0.5f);
                if (processor) {
                    processor->setDownsampleTo31Khz(crush);
                    processor->set16Bit(crush);
                }
                
                break;
            case RobotRocketParameterAddress::compress:
                compress = (value > 0.5f);
                if (processor) {
                    processor->setCompressEnabled(compress);
                }
                
                break;
        }
        
        if (markDirty) {
            storeParam((uint32_t)address, value);
        } else {
            uint32_t addr = (uint32_t)address;
            if (addr < DSPSharedState::kParamCount) {
                uint32_t bits;
                memcpy(&bits, &value, sizeof(float));
                __atomic_store_n(&sharedState.paramValues[addr], bits, __ATOMIC_RELAXED);
            }
        }
    }

    AUAudioFrameCount maximumFramesToRender() const {
        return mMaxFramesToRender;
    }

    void setMaximumFramesToRender(const AUAudioFrameCount &maxFrames) {
        mMaxFramesToRender = maxFrames;
    }

    void process(std::span<float const*> inputBuffers, std::span<float *> outputBuffers, AUEventSampleTime /*bufferStartTime*/, AUAudioFrameCount frameCount) {
        // pass through if bypassed, not licensed or input buffers are missing
        if (bypassed || !licensed || inputBuffers.size() < 2) {
            for (UInt32 channel = 0; channel < std::min(inputBuffers.size(), outputBuffers.size()); ++channel) {
                std::copy_n(inputBuffers[channel], frameCount, outputBuffers[channel]);
            }
            return;
        }

        float* vocalIn = const_cast<float*>(inputBuffers[0]);
        float* synthIn = const_cast<float*>(inputBuffers[1]);
        
        // to keep the vocoder simple it always expects a left and right output, but if the output is mono, just give it the output twice for now
        float* outL = outputBuffers[0];
        float* outR = outputBuffers.size() > 1 ? outputBuffers[1] : outputBuffers[0];

        processor->process(synthIn, vocalIn, outL, outR, static_cast<int>(frameCount));
        
        calculateLevels(synthIn, vocalIn, outL, static_cast<int>(frameCount));
    }
    
    void calculateLevels(float* synthIn, float* voiceIn, float* out, int numFrames) {
        peakMeter->calculateLevels(synthIn, voiceIn, out, numFrames);

        meterSampleCounter += numFrames;

        // periodically report progress back to the UI
        if (meterSampleCounter >= meterUpdateIntervalSamples) {
            storeMeter(peakMeter->getSynthPeakDb(), peakMeter->getVoicePeakDb(), peakMeter->getOutPeakDb());

            meterSampleCounter = 0;
        }
    }

    void handleOneEvent(AUEventSampleTime now, AURenderEvent const *event) {
        switch (event->head.eventType) {
            case AURenderEventParameter: {
                handleParameterEvent(now, event->parameter);
                break;
            }
            default:
                break;
        }
    }

    void handleParameterEvent(AUEventSampleTime /*now*/, AUParameterEvent const& parameterEvent) {
        setParameter(parameterEvent.parameterAddress, parameterEvent.value);
        storeParam((uint32_t)parameterEvent.parameterAddress, parameterEvent.value);
    }

private:
    std::unique_ptr<LPCVocoderProcessor> processor;
    std::unique_ptr<PeakMeter> peakMeter;
    
    double sampleRate = 31250.0;
    
    // configurable parameters
    float quality = 0.0;
    float whiteNoiseLevelDb = 0.0;
    float whiteNoiseThreshold = 0.0;
    float outputGainDb = 0.0;
    float voiceInGainDb = 0.0;
    float synthInGainDb = 0.0;
    int voiceMode = 0;
    float formantShiftCents = 0.0;
    bool highPassFilter = false;
    bool lowPassFilter = false;
    bool voiceFreeze = false;
    bool crush = false;
    bool compress = false;
    
    bool licensed = true;
    bool bypassed = false;
    AUAudioFrameCount mMaxFramesToRender = 1024;
    
    // volume meter variables
    int meterSampleCounter = 0;
    const int meterUpdateIntervalSamples = 48000 / 15; // aim for 15 updates a second
    
    DSPSharedState sharedState;

    void storeParam(uint32_t addr, float value) {
        if (addr < DSPSharedState::kParamCount) {
            uint32_t bits;
            memcpy(&bits, &value, sizeof(float));
            __atomic_store_n(&sharedState.paramValues[addr], bits, __ATOMIC_RELAXED);
            __atomic_fetch_or(&sharedState.dirtyMask, 1u << addr, __ATOMIC_RELEASE);
        }
    }

    void storeMeter(float synth, float voice, float out) {
        storeFloat(sharedState.synthLevel, synth);
        storeFloat(sharedState.voiceLevel, voice);
        storeFloat(sharedState.outLevel,   out);
        __atomic_store_n(&sharedState.meterDirty, 1u, __ATOMIC_RELEASE);
    }

    void storeFloat(uint32_t& field, float value) {
        uint32_t bits;
        memcpy(&bits, &value, sizeof(float));
        __atomic_store_n(&field, bits, __ATOMIC_RELAXED);
    }

    float loadFloat(const uint32_t& field) const {
        uint32_t bits = __atomic_load_n(&field, __ATOMIC_RELAXED);
        float result;
        memcpy(&result, &bits, sizeof(float));
        return result;
    }
};
