// Copyright © 2026 Bjango. All rights reserved.

#pragma once

#import <AudioToolbox/AudioToolbox.h>
#import <AVFoundation/AVFoundation.h>

#include <vector>
#include "DSPKernel.hpp"
#include "BufferedAudioBus.hpp"

class AUProcessHelper
{
public:
    AUProcessHelper(DSPKernel& kernel, BufferedInputBus& mainBus)
    : mKernel{kernel}, mMainInputBus{mainBus} {
    }
    
    void setChannelCounts(UInt32 mainChannelCount, UInt32 outputChannelCount) {
        mMainChannelCount = mainChannelCount;
        mInputBuffers.resize(mainChannelCount);
        mOutputBuffers.resize(outputChannelCount);
    }

    void processWithEvents(AudioBufferList* mainIn, AudioBufferList* outBufferList, AudioTimeStamp const *timestamp, AUAudioFrameCount frameCount, AURenderEvent const *events) {
        AUEventSampleTime now = AUEventSampleTime(timestamp->mSampleTime);
        AUAudioFrameCount framesRemaining = frameCount;
        AURenderEvent const *nextEvent = events;

        auto callProcess = [this, mainIn, outBufferList, frameCount]
                             (AUEventSampleTime now, AUAudioFrameCount framesToDo, AUAudioFrameCount frameOffset) {
                                 
            for (int ch = 0; ch < mMainChannelCount; ++ch) {
                mInputBuffers[ch] = (const float*)mainIn->mBuffers[ch].mData + frameOffset;
            }

            // output channels
            for (int ch = 0; ch < outBufferList->mNumberBuffers; ++ch) {
                mOutputBuffers[ch] = (float*)outBufferList->mBuffers[ch].mData + frameOffset;
            }

            mKernel.process(mInputBuffers, mOutputBuffers, now, framesToDo);
        };
        
        while (framesRemaining > 0) {
            if (!nextEvent) {
                AUAudioFrameCount offset = frameCount - framesRemaining;
                callProcess(now, framesRemaining, offset);
                return;
            }
            auto timeZero = AUEventSampleTime(0);
            auto headTime = nextEvent->head.eventSampleTime;
            AUAudioFrameCount segFrames = AUAudioFrameCount(std::max(timeZero, headTime - now));

            if (segFrames > 0) {
                AUAudioFrameCount offset = frameCount - framesRemaining;
                callProcess(now, segFrames, offset);
                framesRemaining -= segFrames;
                now += AUEventSampleTime(segFrames);
            }

            nextEvent = performAllSimultaneousEvents(now, nextEvent);
        }
    }

    AURenderEvent const * performAllSimultaneousEvents(AUEventSampleTime now, AURenderEvent const *evt) {
        do {
            mKernel.handleOneEvent(now, evt);
            evt = evt->head.next;
        } while (evt && evt->head.eventSampleTime <= now);
        return evt;
    }
    
    AUInternalRenderBlock internalRenderBlock() {
        return ^AUAudioUnitStatus(AudioUnitRenderActionFlags   *actionFlags,
                                  const AudioTimeStamp         *timestamp,
                                  AUAudioFrameCount             frameCount,
                                  NSInteger                     outputBusNumber,
                                  AudioBufferList              *outputData,
                                  const AURenderEvent          *realtimeEvents,
                                  AURenderPullInputBlock pullInputBlock)
        {
            if (frameCount > mKernel.maximumFramesToRender())
                        return kAudioUnitErr_TooManyFramesToProcess;

            AudioUnitRenderActionFlags pullFlags = 0;

            AUAudioUnitStatus err = mMainInputBus.pullInput(&pullFlags, timestamp, frameCount, 0, pullInputBlock);
            if (err) return err;

            AudioBufferList* mainList = mMainInputBus.mutableAudioBufferList;

            // Output pointer validation
            if (!outputData || !outputData->mBuffers[0].mData) {
                return kAudioUnitErr_InvalidProperty;
            }

            processWithEvents(mainList, outputData, timestamp, frameCount, realtimeEvents);
            return noErr;
        };
    }

private:
    DSPKernel& mKernel;
    BufferedInputBus& mMainInputBus;
    std::vector<const float*> mInputBuffers;
    std::vector<float*> mOutputBuffers;
    
    UInt32 mMainChannelCount;
};
