CaptureWrapper.h

// intended to be a cross platform way to do calls to stuff
#pragma once

#include <stdio.h>
#include <iostream>
#include <string>
#include <thread>
#include <vector>
#include <locale>
#include <codecvt>
#include <iostream>

#if WIN32
#include <mmdeviceapi.h>
#include <audiopolicy.h>
#include <audioclient.h>
#include <comdef.h>
#include <Psapi.h>
#include <avrt.h>

#include "./windows/audiocapture.h"
#endif

#if MACOSX
#include "audiocapture.h"
#endif

#include "UsefulTypes.h"

// OO YEA I LOVE PIMPLS
class AudioCapture {
public:
    inline AudioCapture();
    ~AudioCapture() = default;
    
    void registerAudioBufferCallback(callbackFunction functionPointer) {
        impl->registerAudioBufferCallback(functionPointer);
    }
    
    void startAudioCapture() {
        impl->startAudioCapture();
        hasStartedAtLeastOnce = true;
    }

    void stopAudioCapture() {
        if (hasStartedAtLeastOnce) {
            impl->stopAudioCapture();
        }
    }
    
    std::vector<AppStruct> getValidWindows() {
        auto validWindows = impl->getValidWindows();
        
        // Remove the app with the name ${JUCE_APPLICATION_NAME_STRING} from the list
        auto it = std::remove_if(validWindows.begin(), validWindows.end(),
            [](const AppStruct& app) {
                return app.applicationName == JUCE_APPLICATION_NAME_STRING;
            });
        validWindows.erase(it, validWindows.end());

        return validWindows;
    }

    void setActiveProcess(AppStruct app) {
        impl->stopAudioCapture();

        impl->setActiveProcess(app);

        impl->startAudioCapture();
    }

    class ImplBase {
    public:
        // virtual ~ImplBase() = default;
        
        virtual void startAudioCapture() = 0;
        virtual void stopAudioCapture() = 0;

        virtual void registerAudioBufferCallback(callbackFunction functionPointer) = 0;

        virtual std::vector<AppStruct> getValidWindows() = 0;

        virtual void setActiveProcess(AppStruct app) = 0;

        ValidAppsStruct validApps;
    protected:
        callbackFunction audioCallback;

    private:
    };

    std::unique_ptr<ImplBase> impl;

private:
    bool hasStartedAtLeastOnce = false;
};


// specific impl for macos, which uses screencapturekit
#ifdef MACOSX
class MacAudioCapture : public AudioCapture::ImplBase {
public:
    MacAudioCapture() {
        std::cout << "MacAudioCapture constructor" << std::endl;

        currentActiveApp.processID = -1;

        sckAudioCapture.registerAudioBufferCallback(audioCallback);
    }
    ~MacAudioCapture() {
        std::cout << "MacAudioCapture destructor" << std::endl;
    }
    
    void startAudioCapture() {
        std::cout << "MacAudioCapture startAudioCapture" << std::endl;

        readerThread = std::thread([&]() {
            sckAudioCapture.startAudioCapture(currentActiveApp.processID, 1);
        });
    }

    void stopAudioCapture() {
        std::cout << "MacAudioCapture stopAudioCapture" << std::endl;

        sckAudioCapture.stopAudioCapture();

        if (readerThread.joinable())
            readerThread.join();

        // readerThread = std::thread([&]() {
        //     sckAudioCapture.startAudioCapture(currentActiveApp.processID, 1);
        // });

        // sckAudioCapture.startAudioCapture();
    }
    
    void registerAudioBufferCallback(callbackFunction functionPointer) {
        std::cout << "MacAudioCapture registerAudioBufferCallback" << std::endl;

        audioCallback = functionPointer;
        sckAudioCapture.registerAudioBufferCallback(audioCallback);
    }

    std::vector<AppStruct> getValidWindows() {
        std::vector<AppStruct> validWindows;
        auto resp = sckAudioCapture.getValidWindows();

        AppStruct allProcesses;
        allProcesses.processID = -1;
        allProcesses.bundleIdentifier = "allProcesses";
        allProcesses.applicationName = "All Processes";

        validWindows.clear();
        validWindows.push_back(allProcesses);

        for (int i = 1; i < resp.size; i++) {
            validWindows.push_back((AppStruct)resp.apps[i]);
            std::cout << resp.apps[i].applicationName << ", " << resp.apps[i].bundleIdentifier << ", " << resp.apps[i].processID << std::endl;
        }

        return validWindows;
    }

    void setActiveProcess(AppStruct app) {
        currentActiveApp = app;
    }

private:
    SCKAudioCapture sckAudioCapture;

    AppStruct currentActiveApp;

    std::thread readerThread;
};
#endif


// specific impl for windows, which uses wasapi?
// specific impl for macos, which uses screencapturekit
#ifdef WIN32
class WindowsAudioCapture : public AudioCapture::ImplBase {
public:
    WindowsAudioCapture() {
        std::cout << "WindowsAudioCapture constructor" << std::endl;
        
        currentActiveApp.processID = -1;
    }

    ~WindowsAudioCapture() {
        std::cout << "WindowsAudioCapture destructor" << std::endl;
    }
    
    void startAudioCapture() {
        std::cout << "WindowsAudioCapture startAudioCapture" << std::endl;

        if (currentActiveApp.processID == -1) {
            std::cerr << "No active process set, include everything" << std::endl;
            loopbackCapture.StartCaptureAsync(-1, false);
        } else {
            loopbackCapture.StartCaptureAsync(currentActiveApp.processID, true);
        }
    }

    void stopAudioCapture() {
        std::cout << "WindowsAudioCapture stopAudioCapture" << std::endl;

        loopbackCapture.StopCaptureAsync();
    }
    
    void registerAudioBufferCallback(callbackFunction functionPointer) {
        std::cout << "WindowsAudioCapture registerAudioBufferCallback" << std::endl;

        audioCallback = functionPointer;
        loopbackCapture.registerAudioCallback(audioCallback);
    }

    std::vector<AppStruct> getValidWindows() {
        // HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
        // if (FAILED(hr))
        // {
        //     std::cerr << "Failed to initialize COM library." << std::endl;
        //     return std::vector<AppStruct>();
        // }

        IMMDeviceEnumerator *pDeviceEnumerator = nullptr;
        HRESULT hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL, IID_PPV_ARGS(&pDeviceEnumerator));
        if (FAILED(hr))
        {
            std::cerr << "Failed to create device enumerator." << std::endl;
            CoUninitialize();
            return std::vector<AppStruct>();
        }

        IMMDevice *pDefaultDevice = nullptr;
        hr = pDeviceEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &pDefaultDevice);
        if (FAILED(hr))
        {
            std::cerr << "Failed to get default audio endpoint." << std::endl;
            pDeviceEnumerator->Release();
            CoUninitialize();
            return std::vector<AppStruct>();
        }

        IAudioSessionManager2 *pAudioSessionManager = nullptr;
        hr = pDefaultDevice->Activate(__uuidof(IAudioSessionManager2), CLSCTX_ALL, nullptr, (void **)&pAudioSessionManager);
        if (FAILED(hr))
        {
            std::cerr << "Failed to activate audio session manager." << std::endl;
            pDefaultDevice->Release();
            pDeviceEnumerator->Release();
            CoUninitialize();
            return std::vector<AppStruct>();
        }

        IAudioSessionEnumerator *pSessionEnumerator = nullptr;
        hr = pAudioSessionManager->GetSessionEnumerator(&pSessionEnumerator);
        if (FAILED(hr))
        {
            std::cerr << "Failed to get session enumerator." << std::endl;
            pAudioSessionManager->Release();
            pDefaultDevice->Release();
            pDeviceEnumerator->Release();
            CoUninitialize();
            return std::vector<AppStruct>();
        }
        

        int sessionCount = 0;
        hr = pSessionEnumerator->GetCount(&sessionCount);

        validWindows.clear();

        if (SUCCEEDED(hr))
        {
            for (int i = 0; i < sessionCount; ++i)
            {
                IAudioSessionControl *pSessionControl = nullptr;
                hr = pSessionEnumerator->GetSession(i, &pSessionControl);
                if (SUCCEEDED(hr))
                {
                    IAudioSessionControl2 *pSessionControl2;
                    hr = pSessionControl->QueryInterface(IID_PPV_ARGS(&pSessionControl2));
                    if (SUCCEEDED(hr))
                    {
                        DWORD processId = 0;
                        hr = pSessionControl2->GetProcessId(&processId);
                        if (SUCCEEDED(hr))
                        {
                            HANDLE processHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId);
                            if (processHandle)
                            {
                                WCHAR processName[MAX_PATH];
                                if (GetModuleBaseNameW(processHandle, nullptr, processName, MAX_PATH) > 0)
                                {   
                                    AppStruct app;
                                    
                                    std::wstring ws( processName );
                                    std::string appName( ws.begin(), ws.end() );
                                    app.applicationName = appName;
                                    app.processID = processId;

                                    validWindows.push_back(app);
                                }
                                CloseHandle(processHandle);
                            }
                        }
                        pSessionControl2->Release();
                    }
                    pSessionControl->Release();
                }
            }
        }

        AppStruct allProcesses;
        allProcesses.processID = -1;
        allProcesses.bundleIdentifier = "allProcesses";
        allProcesses.applicationName = "All Processes";

        validWindows.push_back(allProcesses);

        pSessionEnumerator->Release();
        pAudioSessionManager->Release();
        pDefaultDevice->Release();
        pDeviceEnumerator->Release();
        CoUninitialize();

        return validWindows;
    }

    void setActiveProcess(AppStruct app) {
        currentActiveApp = app;
    }


private:
    std::vector<AppStruct> validWindows;

    AppStruct currentActiveApp;

    std::thread readerThread;

    CLoopbackCapture loopbackCapture;
};
#endif

AudioCapture::AudioCapture() {
    #if WIN32
    impl = std::make_unique<WindowsAudioCapture>();
    #endif

    #if MACOSX
    impl = std::make_unique<MacAudioCapture>();
    #endif
};