// Copyright © 2025 Bjango. All rights reserved.

#pragma once

#include "LPCvocoder.hpp"
#include "LPCAnalyser.hpp"
#include "AudioUtils.hpp"
#include "VoiceLPCData.hpp"
#include "ModelProcessor.hpp"
#include "SinglePeakMeter.hpp"
#include "ParameterSmoother.hpp"
#import "ParameterAddresses.h"
#import <AudioToolbox/AudioToolbox.h>
#import <algorithm>
#import <vector>
#import <span>
#import <memory>
#include <string>
#include <cstring>

#include "../../../models/Model1Binary.hpp"
#include "../../../models/Model10Binary.hpp"
#include "../../../models/Model11Binary.hpp"

// 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 = 8;
    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 inLevelL = 0;
    uint32_t inLevelR = 0;
    uint32_t outLevelL = 0;
    uint32_t outLevelR = 0;
    uint32_t inHoldL = 0;
    uint32_t inHoldR = 0;
    uint32_t outHoldL = 0;
    uint32_t outHoldR = 0;
    uint32_t isInputStereo = 1;
    uint32_t isOutputStereo = 1;
    uint32_t meterDirty = 0;
};

class DSPKernel {
public:
    static constexpr float kHalfPi = 1.57079632679489661923f;

    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 getAtomicMeterInLevelL()  const { return loadFloat(sharedState.inLevelL); }
    float getAtomicMeterInLevelR()  const { return loadFloat(sharedState.inLevelR); }
    float getAtomicMeterOutLevelL() const { return loadFloat(sharedState.outLevelL); }
    float getAtomicMeterOutLevelR() const { return loadFloat(sharedState.outLevelR); }
    float getAtomicMeterInHoldL()   const { return loadFloat(sharedState.inHoldL); }
    float getAtomicMeterInHoldR()   const { return loadFloat(sharedState.inHoldR); }
    float getAtomicMeterOutHoldL()  const { return loadFloat(sharedState.outHoldL); }
    float getAtomicMeterOutHoldR()  const { return loadFloat(sharedState.outHoldR); }
    bool  getAtomicMeterIsInputStereo()  const { return __atomic_load_n(&sharedState.isInputStereo,  __ATOMIC_RELAXED) != 0; }
    bool  getAtomicMeterIsOutputStereo() const { return __atomic_load_n(&sharedState.isOutputStereo, __ATOMIC_RELAXED) != 0; }
    
    void setupLPCProcessor(LPCVocoderProcessor* proc) {
        proc->setQuality(48);
        proc->setWhiteNoiseLevel(-20);
        proc->setWhiteNoiseThreshold(0.3);
        proc->setOutputGain(0);
        proc->setVoiceInGain(0);
        proc->setSynthInGain(0);
        proc->setVoiceMode(0);
        proc->setFormantShift(0);
        proc->setHighPassFilter(false);
        proc->setLowPassFilter(false);
        proc->setDownsampleTo31Khz(false);
        proc->set16Bit(false);
        proc->setCompressEnabled(false);
        proc->setVoiceOnlyOutput(false);
        proc->setWetMix(1.0f);
        proc->setDryMix(0.0f);
        proc->setFrozenLPC(lpcVoiceData);
    }

    void initialize(double inSampleRate,
                    const VoiceLPCData& voiceAIn, const VoiceLPCData& voiceBIn, const VoiceLPCData& voiceCIn, const VoiceLPCData& voiceDIn) {
        sampleRate = inSampleRate;
        
        voiceA = voiceAIn;
        voiceB = voiceBIn;
        voiceC = voiceCIn;
        voiceD = voiceDIn;
        
        lpcProcessorL = std::make_unique<LPCVocoderProcessor>(inSampleRate);
        lpcProcessorR = std::make_unique<LPCVocoderProcessor>(inSampleRate);

        recalculateLPC();
        
        setupLPCProcessor(lpcProcessorL.get());
        setupLPCProcessor(lpcProcessorR.get());
        
        modelProcessor = std::make_unique<Bjango::ModelProcessor>();
        modelProcessor->setupProcessing(sampleRate, mMaxFramesToRender);
        modelProcessor->setEnabled(false);
        modelProcessor->setInputLimiterThreshold(0.95f);

        dryBufferL.resize(mMaxFramesToRender);
        dryBufferR.resize(mMaxFramesToRender);
        modelCarrierBufferL.resize(mMaxFramesToRender);
        modelCarrierBufferR.resize(mMaxFramesToRender);
        lpcScratchBuffer.resize(mMaxFramesToRender);
        
        inMeterL = std::make_unique<SinglePeakMeter>();
        inMeterR = std::make_unique<SinglePeakMeter>();
        outMeterL = std::make_unique<SinglePeakMeter>();
        outMeterR = std::make_unique<SinglePeakMeter>();
        inMeterL->setSampleRate(static_cast<float>(sampleRate));
        inMeterR->setSampleRate(static_cast<float>(sampleRate));
        outMeterL->setSampleRate(static_cast<float>(sampleRate));
        outMeterR->setSampleRate(static_cast<float>(sampleRate));

        inputGainSmoother.prepare(inSampleRate);
        inputGainSmoother.reset(inputGain);
        outputGainSmoother.prepare(inSampleRate);
        outputGainSmoother.reset(outputGain);
        dryGainSmoother.prepare(inSampleRate);
        dryGainSmoother.reset(std::cos(mix * kHalfPi));
        wetGainSmoother.prepare(inSampleRate);
        wetGainSmoother.reset(std::sin(mix * kHalfPi));
    }

    void deInitialize() {
        lpcProcessorL.reset();
        lpcProcessorR.reset();
        modelProcessor.reset();
    }

    void freeFinishedModels() {
        if (modelProcessor) modelProcessor->freeFinishedModels();
    }

    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 RobotVowelsParameterAddress::voice:
                voice = value;
                recalculateLPC();
                
                break;
            case RobotVowelsParameterAddress::distortion:
                distortion = value;
                updateDistortion(true);
                break;
            case RobotVowelsParameterAddress::mix:
                mix = std::clamp(value * 0.01f, 0.0f, 1.0f);
                dryGainSmoother.setTarget(std::cos(mix * kHalfPi));
                wetGainSmoother.setTarget(std::sin(mix * kHalfPi));
                break;
            case RobotVowelsParameterAddress::vowelMixX:
                vowelMixX = value;
                recalculateLPC();
                
                break;
            case RobotVowelsParameterAddress::vowelMixY:
                vowelMixY = value;
                recalculateLPC();
                
                break;
            case RobotVowelsParameterAddress::processingOrder:
                modelFirst = (value < 0.5f);
                
                break;
            case RobotVowelsParameterAddress::inputGain:
                inputGainSmoother.setTarget(AudioUtils::decibelToScale(value));
                break;
            case RobotVowelsParameterAddress::outputGain:
                outputGainSmoother.setTarget(AudioUtils::decibelToScale(value));
                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 updateDistortion(bool withCrossfade = false) {
        if (!modelProcessor) return;

        if (distortion < 0.5f) {
            if (withCrossfade) modelProcessor->setEnabledWithCrossfade(false);
            else modelProcessor->setEnabled(false);
            return;
        }

        const unsigned char* data = nullptr;
        unsigned int size = 0;
        if (distortion < 1.5f) {
            data = BinaryData::model1_json; size = BinaryData::model1_jsonSize;
        } else if (distortion < 2.5f) {
            data = BinaryData::model11_json; size = BinaryData::model11_jsonSize;
        } else {
            data = BinaryData::model10_json; size = BinaryData::model10_jsonSize;
        }

        if (withCrossfade) modelProcessor->loadModelWithCrossfade(data, size);
        else modelProcessor->loadModel(data, size);
    }
    
    void recalculateLPC() {
        auto voiceData = getCurrentVoice();
        voiceGain = voiceData.getGain();
        
        voiceData.getLPCInterpolationAt(vowelMixX, vowelMixY, lpcVoiceData);
        if (lpcProcessorL) {
            lpcProcessorL->setFrozenLPC(lpcVoiceData);
        }
        if (lpcProcessorR) {
            lpcProcessorR->setFrozenLPC(lpcVoiceData);
        }
    }
    
    VoiceLPCData getCurrentVoice() {
        if (voice < 0.5f) {
            return voiceA;
        } else if (voice < 1.5f) {
            return voiceB;
        } else if (voice < 2.5f) {
            return voiceC;
        } else {
            return voiceD;
        }
    }

    void setMaximumFramesToRender(const AUAudioFrameCount &maxFrames) {
        mMaxFramesToRender = maxFrames;
        dryBufferL.resize(mMaxFramesToRender);
        dryBufferR.resize(mMaxFramesToRender);
        modelCarrierBufferL.resize(mMaxFramesToRender);
        modelCarrierBufferR.resize(mMaxFramesToRender);
        lpcScratchBuffer.resize(mMaxFramesToRender);
    }

    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() < 1) {
            for (UInt32 channel = 0; channel < std::min((int)inputBuffers.size(), (int)outputBuffers.size()); ++channel) {
                std::copy_n(inputBuffers[channel], frameCount, outputBuffers[channel]);
            }
            return;
        }

        int numInChannels = static_cast<int>(inputBuffers.size());
        int numOutChannels = static_cast<int>(outputBuffers.size());
        bool isStereoIn = numInChannels > 1;
        bool isStereoOut = numOutChannels > 1;

        const float* inL = inputBuffers[0];
        const float* inR = isStereoIn ? inputBuffers[1] : inL;

        // Save original dry signal for mixing later
        std::copy_n(inL, frameCount, dryBufferL.begin());
        std::copy_n(inR, frameCount, dryBufferR.begin());
        
        // Prepare gained carrier for first stage processing
        for (int i = 0; i < static_cast<int>(frameCount); ++i) {
            const float ig = inputGainSmoother.process();
            modelCarrierBufferL[i] = inL[i] * ig;
            modelCarrierBufferR[i] = inR[i] * ig;
        }
        
        inMeterL->calculateLevels(modelCarrierBufferL.data(), frameCount);
        inMeterR->calculateLevels(isStereoIn ? modelCarrierBufferR.data() : modelCarrierBufferL.data(), frameCount);

        float* outL = outputBuffers[0];
        float* outR = isStereoOut ? outputBuffers[1] : lpcScratchBuffer.data();

        if (modelFirst) {
            // Stage 1: Model Processor
            const float* modelInputs[2] = { modelCarrierBufferL.data(), modelCarrierBufferR.data() };
            float* modelOutputs[2] = { modelCarrierBufferL.data(), modelCarrierBufferR.data() };
            modelProcessor->processStereo(modelInputs, modelOutputs, static_cast<int>(frameCount));

            // Stage 2: LPC Vocoding
            lpcProcessorL->process(modelCarrierBufferL.data(), outL, lpcScratchBuffer.data(), static_cast<int>(frameCount));
            lpcProcessorR->process(modelCarrierBufferR.data(), lpcScratchBuffer.data(), outR, static_cast<int>(frameCount));
        } else {
            // Stage 1: LPC Vocoding
            lpcProcessorL->process(modelCarrierBufferL.data(), outL, lpcScratchBuffer.data(), static_cast<int>(frameCount));
            lpcProcessorR->process(modelCarrierBufferR.data(), lpcScratchBuffer.data(), outR, static_cast<int>(frameCount));

            // Stage 2: Model Processor
            const float* inputs[2] = { outL, outR };
            float* outputs[2] = { outL, outR };
            modelProcessor->processStereo(inputs, outputs, static_cast<int>(frameCount));
        }

        // dry/wet mix — L and R share one loop so the smoothers advance once per sample
        for (int i = 0; i < static_cast<int>(frameCount); ++i) {
            const float ws = wetGainSmoother.process() * outputGainSmoother.process() * voiceGain;
            const float dg = dryGainSmoother.process();
            outL[i] = (outL[i] * ws) + dryBufferL[i] * dg;
            outR[i] = (outR[i] * ws) + dryBufferR[i] * dg;
        }
        
        outMeterL->calculateLevels(outL, frameCount);
        outMeterR->calculateLevels(isStereoOut ? outR : outL, frameCount);
        
        meterSampleCounter += frameCount;

        // periodically report progress back to the UI
        if (meterSampleCounter >= meterUpdateIntervalSamples) {
            storeMeter(
                inMeterL->getPeakDb(), inMeterR->getPeakDb(), outMeterL->getPeakDb(), outMeterR->getPeakDb(),
                inMeterL->getHoldDb(), inMeterR->getHoldDb(), outMeterL->getHoldDb(), outMeterR->getHoldDb(),
                isStereoIn, isStereoOut);

            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, true);
    }

private:
    std::unique_ptr<LPCVocoderProcessor> lpcProcessorL;
    std::unique_ptr<LPCVocoderProcessor> lpcProcessorR;
    std::unique_ptr<Bjango::ModelProcessor> modelProcessor;
    
    // volume meter variables
    int meterSampleCounter = 0;
    const int meterUpdateIntervalSamples = 48000 / 15; // aim for 15 updates a second
    
    std::unique_ptr<SinglePeakMeter> inMeterL;
    std::unique_ptr<SinglePeakMeter> inMeterR;
    std::unique_ptr<SinglePeakMeter> outMeterL;
    std::unique_ptr<SinglePeakMeter> outMeterR;
    
    double sampleRate = 44100.0;
    
    float voiceGain = 1.0f;
    
    // configurable parameters
    float voice = 0.0f;
    float distortion = 0.0f;
    float mix = 0.5f;
    float vowelMixX = 0.0f;
    float vowelMixY = 0.0f;

    // initial scale values — used only to reset smoothers in initialize()
    float inputGain = 1.0f;
    float outputGain = AudioUtils::decibelToScale(3.0f);

    ParameterSmoother inputGainSmoother;
    ParameterSmoother outputGainSmoother;
    ParameterSmoother dryGainSmoother;
    ParameterSmoother wetGainSmoother;
    
    VoiceLPCData voiceA;
    VoiceLPCData voiceB;
    VoiceLPCData voiceC;
    VoiceLPCData voiceD;
    
    LPCAnalysisResult lpcVoiceData;
    
    bool licensed = true;
    bool bypassed = false;
    bool modelFirst = false;
    AUAudioFrameCount mMaxFramesToRender = 1024;
    std::vector<float> dryBufferL;
    std::vector<float> dryBufferR;
    std::vector<float> modelCarrierBufferL;
    std::vector<float> modelCarrierBufferR;
    std::vector<float> lpcScratchBuffer;

    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 inL, float inR, float outL, float outR, float inHL, float inHR, float outHL, float outHR, bool isStereoIn, bool isStereoOut) {
        storeFloat(sharedState.inLevelL,  inL);
        storeFloat(sharedState.inLevelR,  inR);
        storeFloat(sharedState.outLevelL, outL);
        storeFloat(sharedState.outLevelR, outR);
        storeFloat(sharedState.inHoldL,   inHL);
        storeFloat(sharedState.inHoldR,   inHR);
        storeFloat(sharedState.outHoldL,  outHL);
        storeFloat(sharedState.outHoldR,  outHR);
        __atomic_store_n(&sharedState.isInputStereo,  (uint32_t)isStereoIn,  __ATOMIC_RELAXED);
        __atomic_store_n(&sharedState.isOutputStereo, (uint32_t)isStereoOut, __ATOMIC_RELAXED);
        __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;
    }
};
