// livesub_capture.exe — 특정 프로세스(와 그 자식 프로세스)의 렌더 오디오만
// 뽑아내 stdout 으로 흘려보내는 보조 프로그램.
//
// 출력 포맷: 16000Hz / mono / 32-bit float, 헤더 없는 생 PCM.
// 사용법:  livesub_capture.exe --pid 1234 [--exclude]
//          --exclude 를 주면 해당 프로세스 트리를 "제외한" 나머지 소리를 받는다.
//
// Windows 10 2004(build 19041) 이상에서 동작하는 WASAPI process loopback API를
// 사용한다. Microsoft 의 ApplicationLoopback 샘플이 기반이다.

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#include <audioclient.h>
#include <audioclientactivationparams.h>
#include <mmdeviceapi.h>
#include <wrl/implements.h>

#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <io.h>
#include <vector>

#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "mmdevapi.lib")

using Microsoft::WRL::ComPtr;
using Microsoft::WRL::RuntimeClass;
using Microsoft::WRL::RuntimeClassFlags;
using Microsoft::WRL::ClassicCom;
using Microsoft::WRL::FtmBase;

namespace {

constexpr int kSampleRate = 16000;
constexpr int kChannels = 1;

std::atomic<bool> g_stop{false};

BOOL WINAPI ConsoleHandler(DWORD type) {
    if (type == CTRL_C_EVENT || type == CTRL_BREAK_EVENT || type == CTRL_CLOSE_EVENT) {
        g_stop = true;
        return TRUE;
    }
    return FALSE;
}

// ActivateAudioInterfaceAsync 완료 콜백. 결과를 이벤트로 넘겨준다.
class ActivationHandler
    : public RuntimeClass<RuntimeClassFlags<ClassicCom | FtmBase>,
                          IActivateAudioInterfaceCompletionHandler> {
public:
    explicit ActivationHandler(HANDLE done) : done_(done) {}

    STDMETHODIMP ActivateCompleted(IActivateAudioInterfaceAsyncOperation* op) override {
        HRESULT activate_hr = S_OK;
        ComPtr<IUnknown> unknown;
        HRESULT hr = op->GetActivateResult(&activate_hr, &unknown);
        if (SUCCEEDED(hr) && SUCCEEDED(activate_hr)) {
            unknown.As(&client_);
        }
        result_ = FAILED(hr) ? hr : activate_hr;
        SetEvent(done_);
        return S_OK;
    }

    HRESULT result() const { return result_; }
    ComPtr<IAudioClient> client() const { return client_; }

private:
    HANDLE done_;
    HRESULT result_ = E_FAIL;
    ComPtr<IAudioClient> client_;
};

void Fail(const char* what, HRESULT hr) {
    std::fprintf(stderr, "[livesub_capture] %s (hr=0x%08lX)\n", what, static_cast<unsigned long>(hr));
}

int Run(DWORD pid, bool exclude) {
    HANDLE activated = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    HANDLE sample_ready = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    if (!activated || !sample_ready) {
        Fail("이벤트 생성 실패", HRESULT_FROM_WIN32(GetLastError()));
        return 1;
    }

    AUDIOCLIENT_ACTIVATION_PARAMS params{};
    params.ActivationType = AUDIOCLIENT_ACTIVATION_TYPE_PROCESS_LOOPBACK;
    params.ProcessLoopbackParams.TargetProcessId = pid;
    params.ProcessLoopbackParams.ProcessLoopbackMode =
        exclude ? PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE
                : PROCESS_LOOPBACK_MODE_INCLUDE_TARGET_PROCESS_TREE;

    PROPVARIANT activate_params{};
    activate_params.vt = VT_BLOB;
    activate_params.blob.cbSize = sizeof(params);
    activate_params.blob.pBlobData = reinterpret_cast<BYTE*>(&params);

    auto handler = Microsoft::WRL::Make<ActivationHandler>(activated);
    ComPtr<IActivateAudioInterfaceAsyncOperation> async_op;
    HRESULT hr = ActivateAudioInterfaceAsync(VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK,
                                             __uuidof(IAudioClient), &activate_params,
                                             handler.Get(), &async_op);
    if (FAILED(hr)) {
        Fail("ActivateAudioInterfaceAsync 실패", hr);
        return 2;
    }
    if (WaitForSingleObject(activated, 5000) != WAIT_OBJECT_0 || FAILED(handler->result())) {
        Fail("오디오 인터페이스 활성화 실패 (Windows 10 2004 이상 필요)", handler->result());
        return 3;
    }

    ComPtr<IAudioClient> client = handler->client();
    if (!client) {
        Fail("IAudioClient 를 받지 못했습니다", E_POINTER);
        return 3;
    }

    // process loopback 경로에는 GetMixFormat 이 없으므로 포맷을 직접 지정한다.
    WAVEFORMATEX format{};
    format.wFormatTag = WAVE_FORMAT_IEEE_FLOAT;
    format.nChannels = kChannels;
    format.nSamplesPerSec = kSampleRate;
    format.wBitsPerSample = 32;
    format.nBlockAlign = format.nChannels * format.wBitsPerSample / 8;
    format.nAvgBytesPerSec = format.nSamplesPerSec * format.nBlockAlign;
    format.cbSize = 0;

    constexpr REFERENCE_TIME kBufferDuration = 2 * 10'000'000LL;  // 200ms
    hr = client->Initialize(AUDCLNT_SHAREMODE_SHARED,
                            AUDCLNT_STREAMFLAGS_LOOPBACK | AUDCLNT_STREAMFLAGS_EVENTCALLBACK,
                            kBufferDuration, 0, &format, nullptr);
    if (FAILED(hr)) {
        Fail("IAudioClient::Initialize 실패", hr);
        return 4;
    }

    ComPtr<IAudioCaptureClient> capture;
    hr = client->GetService(IID_PPV_ARGS(&capture));
    if (FAILED(hr)) {
        Fail("IAudioCaptureClient 획득 실패", hr);
        return 4;
    }
    hr = client->SetEventHandle(sample_ready);
    if (FAILED(hr)) {
        Fail("SetEventHandle 실패", hr);
        return 4;
    }
    hr = client->Start();
    if (FAILED(hr)) {
        Fail("스트림 시작 실패", hr);
        return 4;
    }

    std::fprintf(stderr, "[livesub_capture] pid=%lu 캡처 시작 (%dHz mono f32)\n",
                 static_cast<unsigned long>(pid), kSampleRate);
    std::fflush(stderr);

    while (!g_stop) {
        // 무음 구간에서는 이벤트가 오지 않으므로 타임아웃으로 빠져나와 종료 여부를 확인한다.
        DWORD wait = WaitForSingleObject(sample_ready, 200);
        if (wait == WAIT_TIMEOUT) continue;
        if (wait != WAIT_OBJECT_0) break;

        UINT32 packet = 0;
        while (SUCCEEDED(capture->GetNextPacketSize(&packet)) && packet > 0) {
            BYTE* data = nullptr;
            UINT32 frames = 0;
            DWORD flags = 0;
            hr = capture->GetBuffer(&data, &frames, &flags, nullptr, nullptr);
            if (FAILED(hr)) break;

            const size_t bytes = static_cast<size_t>(frames) * format.nBlockAlign;
            if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
                static thread_local std::vector<BYTE> zeros;
                zeros.assign(bytes, 0);
                std::fwrite(zeros.data(), 1, bytes, stdout);
            } else if (data && bytes) {
                std::fwrite(data, 1, bytes, stdout);
            }
            std::fflush(stdout);
            capture->ReleaseBuffer(frames);

            if (std::ferror(stdout)) {  // 부모 프로세스가 파이프를 닫음
                g_stop = true;
                break;
            }
        }
    }

    client->Stop();
    CloseHandle(activated);
    CloseHandle(sample_ready);
    return 0;
}

}  // namespace

int main(int argc, char** argv) {
    DWORD pid = 0;
    bool exclude = false;
    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--pid") == 0 && i + 1 < argc) {
            pid = static_cast<DWORD>(std::strtoul(argv[++i], nullptr, 10));
        } else if (std::strcmp(argv[i], "--exclude") == 0) {
            exclude = true;
        }
    }
    if (pid == 0) {
        std::fprintf(stderr, "사용법: livesub_capture.exe --pid <PID> [--exclude]\n");
        return 64;
    }

    _setmode(_fileno(stdout), _O_BINARY);
    SetConsoleCtrlHandler(ConsoleHandler, TRUE);

    HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (FAILED(hr)) {
        Fail("CoInitializeEx 실패", hr);
        return 1;
    }
    int rc = Run(pid, exclude);
    CoUninitialize();
    return rc;
}