#include <iostream>
#include <windows.h>
#include <signal.h>
#include "LoopbackCapture.h" // Assuming this is where CLoopbackCapture is defined

void usage() {
    std::wcout << L"Usage: AudioCapture <ProcessId> <includetree|excludetree> <OutputFile>\n";
}

volatile bool keepRunning = true;

void signalHandler(int signum) {
    keepRunning = false;
}

int wmain(int argc, wchar_t* argv[])
{
    if (argc != 4)
    {
        usage();
        return 0;
    }

    DWORD processId = wcstoul(argv[1], nullptr, 0);
    if (processId == 0)
    {
        usage();
        return 0;
    }

    bool includeProcessTree;
    if (wcscmp(argv[2], L"includetree") == 0)
    {
        includeProcessTree = true;
    }
    else if (wcscmp(argv[2], L"excludetree") == 0)
    {
        includeProcessTree = false;
    }
    else
    {
        usage();
        return 0;
    }

    PCWSTR outputFile = argv[3];

    // Set up signal handling
    signal(SIGINT, signalHandler);
    signal(SIGTERM, signalHandler);

    CLoopbackCapture loopbackCapture;
    HRESULT hr = loopbackCapture.StartCaptureAsync(processId, includeProcessTree, outputFile);
    if (FAILED(hr))
    {
        wil::unique_hlocal_string message;
        FormatMessageW(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS | FORMAT_MESSAGE_ALLOCATE_BUFFER, nullptr, hr,
            MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (PWSTR)&message, 0, nullptr);
        std::wcout << L"Failed to start capture\n0x" << std::hex << hr << L": " << message.get() << L"\n";
    }
    else
    {
        std::wcout << L"Capturing audio. Terminate the process to stop." << std::endl;

        // Run until a termination signal is received
        while (keepRunning)
        {
            Sleep(100); // Sleep to reduce CPU usage
        }

        // Stop the capture
        loopbackCapture.StopCaptureAsync();

        std::wcout << L"Capture stopped." << std::endl;
    }

    return 0;
}