CoopAllTheThings/hook/src/audio_hook.cpp

#include "audio_hook.hpp"

#include <atomic>
#include <cstring>
#include <mutex>
#include <unordered_map>

#include <windows.h>

#include <audioclient.h>
#include <mmdeviceapi.h>
#include <mmreg.h>

#include "debug_log.hpp"
#include "hook_registry.hpp"
#include "rate_estimator.hpp"

namespace coop::hook
{

namespace
{

// COM vtable indices (frozen ABI). IUnknown occupies 0..2.
//   IMMDevice:          Activate = 3
//   IAudioClient:       Initialize = 3, ... SetEventHandle = 13, GetService = 14
//   IAudioRenderClient: GetBuffer = 3, ReleaseBuffer = 4
constexpr unsigned kIdx_IMMDevice_Activate = 3;
constexpr unsigned kIdx_IAudioClient_Initialize = 3;
constexpr unsigned kIdx_IAudioClient_GetService = 14;
constexpr unsigned kIdx_IAudioRenderClient_GetBuffer = 3;
constexpr unsigned kIdx_IAudioRenderClient_ReleaseBuffer = 4;

// Original COM method signatures (all __stdcall via STDMETHODCALLTYPE). We call
// the originals through the saved vtable pointers, so these types must match the
// real interfaces exactly.
using ActivateFn = HRESULT(STDMETHODCALLTYPE*)(IMMDevice*, REFIID, DWORD, PROPVARIANT*, void**);
using InitializeFn = HRESULT(STDMETHODCALLTYPE*)(IAudioClient*, AUDCLNT_SHAREMODE, DWORD, REFERENCE_TIME,
                                                 REFERENCE_TIME, const WAVEFORMATEX*, LPCGUID);
using GetServiceFn = HRESULT(STDMETHODCALLTYPE*)(IAudioClient*, REFIID, void**);
using GetBufferFn = HRESULT(STDMETHODCALLTYPE*)(IAudioRenderClient*, UINT32, BYTE**);
using ReleaseBufferFn = HRESULT(STDMETHODCALLTYPE*)(IAudioRenderClient*, UINT32, DWORD);

// Hooks one COM vtable slot by overwriting its function pointer; the original is
// called through the saved pointer. We use this instead of SafetyHook's inline
// hooks for the WASAPI COM methods because, on x86, MMDevApi/AudioSes prologues
// use dynamic stack alignment (`and esp,-8`) with EBP-relative argument access,
// which SafetyHook's trampoline relocation mishandles: the relocated prologue
// leaves EBP wrong, so the original reads garbage arguments and faults (it froze
// 32-bit FMOD games the instant audio init ran through the hook). Swapping the
// vtable entry leaves the original code untouched, so it runs with a pristine
// stack regardless of prologue shape. Every instance of a COM coclass shares one
// vtable, so a single swap intercepts all of them (the same property the old
// inline approach relied on). See the project's stdcall-x86 note.
class VtableHook
{
public:
	bool install(void* com_object, unsigned index, void* detour)
	{
		if (m_vtable != nullptr)
		{
			return true; // already installed (shared vtable covers every instance)
		}
		auto** vtable = *reinterpret_cast<void***>(com_object);
		DWORD old_protect = 0;
		if (!VirtualProtect(&vtable[index], sizeof(void*), PAGE_READWRITE, &old_protect))
		{
			return false;
		}
		m_original = vtable[index];
		vtable[index] = detour; // aligned pointer store -> atomic vs. a concurrent caller
		VirtualProtect(&vtable[index], sizeof(void*), old_protect, &old_protect);
		m_vtable = vtable;
		m_index = index;
		return true;
	}

	void remove()
	{
		if (m_vtable == nullptr)
		{
			return;
		}
		DWORD old_protect = 0;
		if (VirtualProtect(&m_vtable[m_index], sizeof(void*), PAGE_READWRITE, &old_protect))
		{
			m_vtable[m_index] = m_original;
			VirtualProtect(&m_vtable[m_index], sizeof(void*), old_protect, &old_protect);
		}
		m_vtable = nullptr;
		// Deliberately keep m_original valid: a detour already running on the audio thread
		// (it doesn't hold our setup lock) may still call original() after we restore the
		// slot. The original function lives in the loaded audio module, so the pointer stays
		// valid; nulling it here would race that in-flight detour into a null call (a rapid
		// hook/unhook crash the mock-game stress test caught). A re-install re-reads it.
		m_index = 0;
	}

	template <typename Fn> Fn original() const { return reinterpret_cast<Fn>(m_original); }
	explicit operator bool() const { return m_vtable != nullptr; }

private:
	void** m_vtable = nullptr;
	unsigned m_index = 0;
	void* m_original = nullptr;
};

// The scalar audio format we forward; resolved from the game's WAVEFORMATEX.
struct CapturedFormat
{
	std::uint32_t rate = 0;
	std::uint32_t channels = 0;
	std::uint32_t bits = 0;
	std::uint32_t tag = 0; // WAVE_FORMAT_PCM / _IEEE_FLOAT (EXTENSIBLE resolved to its subformat)
	std::uint32_t block_align = 0;
};

// --- Global hook state -----------------------------------------------------

// Atomic: the audio hot path (hk_ReleaseBuffer) reads it without the setup lock, while
// remove_audio_hooks nulls it under the lock during an unhook. A plain pointer was a
// TOCTOU null-deref (check non-null, then it's nulled, then the call) -- a rapid
// hook/unhook crash the mock-game stress test caught. Load it once and use that.
std::atomic<IpcClient*> g_ipc{nullptr};
// One ring per tracked stream (index = the stream's debug slot). Stream 0 is the
// primary; the host creates a ring per stream and mixes them.
std::atomic<AudioRingHeader*> g_rings[kMaxAudioStreams]{};

std::mutex g_setup_mutex; // guards installs + the format map + stream registration

VtableHook g_vh_activate;
VtableHook g_vh_initialize;
VtableHook g_vh_getservice;
VtableHook g_vh_getbuffer;
VtableHook g_vh_releasebuffer;
bool g_audioclient_hooked = false;
// Bumped on every detour install/remove. hk_GetBuffer stamps the current epoch into a
// thread-local; hk_ReleaseBuffer captures only if the epoch still matches -- so a
// GetBuffer/ReleaseBuffer pair that straddles a hook toggle (the stashed buffer pointer is
// then stale) is skipped instead of memset-ing a freed buffer (a rapid hook/unhook crash
// the mock-game stress test caught).
std::atomic<std::uint32_t> g_hook_epoch{0};
// Number of detours (hk_GetBuffer/hk_ReleaseBuffer) currently executing on the audio
// thread. remove_audio_hooks restores the vtable slots (so no NEW detour starts) and then
// waits for this to drain to 0 before tearing down shared state -- guaranteeing no detour
// is mid-flight when state is cleared. The classic safe-unhook race; the alternative was an
// intermittent access violation in the game during a hook/unhook (the stress test caught it).
std::atomic<int> g_detours_active{0};
struct DetourGuard
{
	DetourGuard()
	{
		g_detours_active.fetch_add(1, std::memory_order_acq_rel);
	}
	~DetourGuard()
	{
		g_detours_active.fetch_sub(1, std::memory_order_acq_rel);
	}
};

// Registry ids for the hook list.
int g_id_activate = -1;
int g_id_initialize = -1;
int g_id_getservice = -1;
int g_id_getbuffer = -1;
int g_id_releasebuffer = -1;

// Our own probe COM objects, created at anchor time purely to read the shared
// IAudioClient / IAudioRenderClient vtables and hook GetBuffer/ReleaseBuffer
// *proactively* — so render clients the game created before we injected (the
// common case: we attach to a game that's already playing) are still caught.
// g_self_render is excluded from capture (it's silent and never rendered).
IMMDevice* g_self_device = nullptr; // kept alive too, so the Activate hook can be re-installed
IAudioClient* g_self_client = nullptr;
std::atomic<IAudioRenderClient*> g_self_render{nullptr};
// The probe objects above are built ONCE and kept for the DLL's lifetime; an audio
// enable/disable toggle then only swaps vtable slots, never creates/destroys COM objects.
// Rapid create/destroy raced AudioSes and crashed the game (the mock-game stress test
// caught this). Released only on detach (shutdown_audio_hooks).

// Device mix format, captured from our probe client. Used as the assumed format
// for a stream we discover on the hot path (we never saw its Initialize, so we
// can't know its real format; shared-mode clients overwhelmingly use the mix
// format). Written once at install before any hook is live.
CapturedFormat g_mix_format;
std::atomic<std::uint32_t> g_have_mix_format{0};

// Each tracked stream's actual format, captured when it's registered. The host
// creates the rings only when audio mirroring is toggled on — typically *after* a
// stream was already registered — so the format must be (re)published to a ring
// whenever it attaches. Guarded by g_setup_mutex.
CapturedFormat g_stream_formats[kMaxAudioStreams];

// Per IAudioClient, the format captured at Initialize, looked up when its render
// client is created. Setup-path only (never touched on the audio thread).
std::unordered_map<IAudioClient*, CapturedFormat> g_client_formats;

// Streams we track (frame counting + per-stream capture). Index 0 is primary.
struct TrackedStream
{
	std::atomic<IAudioRenderClient*> client{nullptr};
	std::atomic<std::uint64_t> frames{0};
	std::atomic<std::uint32_t> block_align{0};     // hot-path frame size for this stream
	std::atomic<std::uint32_t> assumed_format{0};  // 1 = channels/bits guessed -> clamp copies safely
};
TrackedStream g_streams[kMaxAudioStreams];
std::uint32_t g_registered = 0;            // slots filled (<= kMaxAudioStreams), under mutex
std::atomic<std::uint32_t> g_streams_seen{0}; // total distinct clients ever seen

std::atomic<std::uint64_t> g_frames_captured{0}; // total frames captured across streams
std::atomic<std::uint64_t> g_frames_silenced{0}; // total frames whose local playback we muted (SILENT)

// When we attach to an already-running game we never saw its IAudioClient::Initialize,
// so a render client discovered on the hot path gets the device mix format as a best
// guess. That guess is wrong for games that render at a non-device rate via WASAPI
// AUTOCONVERTPCM -- e.g. Godot/Brotato render 44100 while the device mixes at 48000, so
// playing the captured 44100 audio back as 48000 shifts the pitch up. For such streams we
// verify (and correct) the guessed sample rate by measuring the real render cadence
// before publishing the format. g_stream_rate_guess marks a guessed stream; g_rate_estimator
// is its robust, consensus-based measurement (see rate_estimator.hpp). Both guarded by
// g_setup_mutex.
bool g_stream_rate_guess[kMaxAudioStreams] = {};
RateEstimator g_rate_estimator[kMaxAudioStreams] = {};

// We can measure a guessed stream's sample rate, but channels/bits aren't recoverable for a
// client we never saw Initialize -- they stay the device-mix guess. That guess is right for
// the common case (games render stereo float, matching the device, just at a different
// rate), but if a game renders a different channel/bit layout the guessed bytes-per-frame is
// too large and the capture copy would over-read the game's buffer. We can't detect the true
// layout (AUTOCONVERTPCM hands back a fixed staging buffer, so there's no buffer-stride to
// measure, and WASAPI exposes no API for a pre-existing client's format), so we instead clamp
// every guessed-stream copy to the source buffer's committed region (copy_bound_locked /
// readable_bytes) -- the audio may be misinterpreted, but it can never read past the
// allocation. See README "Lessons learned".

// Per-stream AudioFormatState (how its format was determined), mirrored to the host UI.
std::uint32_t g_stream_format_state[kMaxAudioStreams] = {};

// Last operator-op sequence applied per stream (host posts re-measure / override via the
// ring's op channel; we apply each new op once). Guarded by g_setup_mutex.
std::uint32_t g_last_op_seq[kMaxAudioStreams] = {};

// GetBuffer/ReleaseBuffer are paired on one thread, never nested: stash the
// pointer the game just got so ReleaseBuffer can copy it before releasing.
thread_local IAudioRenderClient* t_gb_client = nullptr;
thread_local BYTE* t_gb_data = nullptr;
thread_local UINT32 t_gb_frames = 0;
thread_local std::uint32_t t_gb_epoch = 0; // hook epoch at the GetBuffer (see g_hook_epoch)

CapturedFormat capture_format(const WAVEFORMATEX* wfx)
{
	CapturedFormat cf;
	cf.rate = wfx->nSamplesPerSec;
	cf.channels = wfx->nChannels;
	cf.bits = wfx->wBitsPerSample;
	cf.block_align = wfx->nBlockAlign;
	cf.tag = wfx->wFormatTag;
	if (wfx->wFormatTag == WAVE_FORMAT_EXTENSIBLE && wfx->cbSize >= 22)
	{
		const auto* ext = reinterpret_cast<const WAVEFORMATEXTENSIBLE*>(wfx);
		if (ext->SubFormat == KSDATAFORMAT_SUBTYPE_IEEE_FLOAT)
		{
			cf.tag = WAVE_FORMAT_IEEE_FLOAT;
		}
		else if (ext->SubFormat == KSDATAFORMAT_SUBTYPE_PCM)
		{
			cf.tag = WAVE_FORMAT_PCM;
		}
	}
	if (cf.block_align == 0)
	{
		cf.block_align = cf.channels * (cf.bits / 8);
	}
	return cf;
}

// --- Detours (declared before the installers that reference them) -----------

bool stream_tracked(IAudioRenderClient* rc);
void try_register_lazy(IAudioRenderClient* rc);

// Bytes safely readable from `ptr` within its committed region. Used to cap a guessed
// stream's copy: if its channels/bits differ from the device guess the assumed block is too
// large, and this stops the capture copy from reading past the source buffer's allocation
// (the data is then misinterpreted, but it can never AV). A no-op when the guess is right.
std::uint32_t readable_bytes(const void* ptr, std::uint32_t want)
{
	MEMORY_BASIC_INFORMATION mbi{};
	if (VirtualQuery(ptr, &mbi, sizeof(mbi)) == sizeof(mbi) && mbi.State == MEM_COMMIT)
	{
		const auto* base = static_cast<const std::uint8_t*>(mbi.BaseAddress);
		const auto avail = static_cast<std::uintptr_t>((base + mbi.RegionSize) -
		                                               static_cast<const std::uint8_t*>(ptr));
		if (avail < want)
		{
			return static_cast<std::uint32_t>(avail);
		}
	}
	return want;
}

HRESULT STDMETHODCALLTYPE hk_GetBuffer(IAudioRenderClient* self, UINT32 num_frames, BYTE** data)
{
	DetourGuard guard; // counts this detour as in-flight (drained before an unhook tears down)
	hook_note_call(g_id_getbuffer);
	const HRESULT hr = g_vh_getbuffer.original<GetBufferFn>()(self, num_frames, data);
	if (SUCCEEDED(hr) && data != nullptr)
	{
		t_gb_client = self;
		t_gb_data = *data;
		t_gb_frames = num_frames;
		t_gb_epoch = g_hook_epoch.load(std::memory_order_acquire);
	}
	return hr;
}

HRESULT STDMETHODCALLTYPE hk_ReleaseBuffer(IAudioRenderClient* self, UINT32 num_frames, DWORD flags)
{
	DetourGuard guard; // counts this detour as in-flight (drained before an unhook tears down)
	hook_note_call(g_id_releasebuffer);
	// A render client we've never seen actively rendering is almost certainly one
	// the game created before we injected; adopt it now (the first becomes the
	// primary we capture). Skip our own silent probe client.
	if (self != g_self_render.load(std::memory_order_acquire) && num_frames > 0 && !stream_tracked(self))
	{
		try_register_lazy(self);
	}

	// Per tracked stream: count frames (debug view) and, into the stream's own ring,
	// capture + silence its buffer while capture is enabled. Every stream is captured
	// into its own ring; the host mixes them.
	for (std::uint32_t i = 0; i < kMaxAudioStreams; ++i)
	{
		if (g_streams[i].client.load(std::memory_order_acquire) != self)
		{
			continue;
		}
		const std::uint64_t total =
		    g_streams[i].frames.fetch_add(num_frames, std::memory_order_relaxed) + num_frames;
		if (IpcClient* ipc = g_ipc.load(std::memory_order_acquire)) // load once (unhook may null it)
		{
			ipc->note_audio_frames(i, total);
		}

		if (num_frames > 0 && (flags & AUDCLNT_BUFFERFLAGS_SILENT) == 0)
		{
			AudioRingHeader* ring = g_rings[i].load(std::memory_order_acquire);
			// Only capture once the format is published -- for a guessed-rate stream that's
			// after the true rate is measured, so we never capture/silence audio we'd
			// mis-rate (and don't build a backlog while measuring; the game stays audible).
			if (ring != nullptr && ring->capture_enabled.load(std::memory_order_relaxed) != 0 &&
			    audio_ring_format_ready(*ring) && t_gb_client == self && t_gb_data != nullptr &&
			    t_gb_frames == num_frames &&
			    t_gb_epoch == g_hook_epoch.load(std::memory_order_acquire)) // same hooked epoch as the GetBuffer
			{
				const bool guessed = g_streams[i].assumed_format.load(std::memory_order_relaxed) != 0;
				const std::uint32_t block = g_streams[i].block_align.load(std::memory_order_relaxed);
				std::uint32_t bytes = num_frames * block;
				// A guessed (pre-existing client) stream's block may be larger than the real
				// per-frame size, so clamp the COPY to what's actually readable -- never over-read
				// the game's buffer (no-op when the guess is right).
				if (guessed)
				{
					bytes = readable_bytes(t_gb_data, bytes);
				}
				// Only silence if the frames made it into the ring; if the host has
				// stalled (ring full) keep playing locally rather than going dead
				// silent — degrades to today's echo, never to silence.
				if (block != 0 && audio_ring_push(*ring, t_gb_data, bytes, num_frames))
				{
					g_frames_captured.fetch_add(num_frames, std::memory_order_relaxed);
					// Mute the game's local playback so the only audio is the host's re-render.
					// Otherwise the game plays locally AND the mirror re-renders the same audio a
					// few ms later = a metallic double (the Brotato symptom). AUDCLNT_BUFFERFLAGS_SILENT
					// tells WASAPI to treat the buffer as silence and IGNORE its contents, so it
					// mutes WITHOUT writing the buffer -- safe even for a guessed-format stream whose
					// true frame size we don't know. (Muting used to be tied to the memset below,
					// which is unsafe for a guessed block, so guessed streams -- the late-attach /
					// Brotato case -- were captured but left audible. The flag is the actual mute;
					// the memset is not needed for it.) For an exact/override format we additionally
					// zero the buffer (belt-and-suspenders; `block` is the real frame size there, so
					// it stays in-bounds). Only mutes once the frames made the ring (above) -- a
					// stalled host degrades to echo, never to dead silence.
					if (!guessed)
					{
						std::memset(t_gb_data, 0, bytes);
					}
					g_frames_silenced.fetch_add(num_frames, std::memory_order_relaxed);
					return g_vh_releasebuffer.original<ReleaseBufferFn>()(
					    self, num_frames, flags | AUDCLNT_BUFFERFLAGS_SILENT);
				}
			}
		}
		// Format-verification co-capture (host-driven, step 2a). While the host has set
		// verify_capture and this guessed stream's rate is still being MEASURED (format not yet
		// published), push the raw pre-mix bytes to the ring WITHOUT silencing, so the host can
		// capture both the hook (pre-mix) and a parallel process-loopback (post-mix) of the same
		// audio and cross-correlate them to recover the true format from ground truth. The game
		// stays audible (the host runs loopback during the measurement window anyway), and the
		// shipping no-echo capture/silence path above is left completely untouched.
		if (num_frames > 0 && (flags & AUDCLNT_BUFFERFLAGS_SILENT) == 0)
		{
			AudioRingHeader* vring = g_rings[i].load(std::memory_order_acquire);
			if (vring != nullptr && vring->verify_capture.load(std::memory_order_relaxed) != 0 &&
			    !audio_ring_format_ready(*vring) &&
			    g_streams[i].assumed_format.load(std::memory_order_relaxed) != 0 && t_gb_client == self &&
			    t_gb_data != nullptr && t_gb_frames == num_frames &&
			    t_gb_epoch == g_hook_epoch.load(std::memory_order_acquire))
			{
				const std::uint32_t block = g_streams[i].block_align.load(std::memory_order_relaxed);
				if (block != 0)
				{
					// Self-describing chunk: [u32 frame-count][num_frames*block bytes]. The host can't
					// know the real frame size of a guessed stream, so it recovers the layout by
					// trying candidate de-interleavings -- but it needs the frame count to strip the
					// per-buffer padding (the guessed/device block over-reads a stream with fewer
					// channels/bits). Push both parts only if both fit and the payload is fully
					// readable, so a full ring or a short buffer can never tear the framing.
					const std::uint32_t want = num_frames * block;
					if (readable_bytes(t_gb_data, want) == want &&
					    audio_ring_free_space(*vring) >= static_cast<std::uint32_t>(sizeof(num_frames)) + want)
					{
						audio_ring_push(*vring, &num_frames, sizeof(num_frames), 0);
						audio_ring_push(*vring, t_gb_data, want, num_frames);
					}
				}
			}
		}
		break;
	}
	return g_vh_releasebuffer.original<ReleaseBufferFn>()(self, num_frames, flags);
}

// Publish a stream's format + state to the host's per-stream debug channel. Caller holds
// g_setup_mutex.
void publish_stream_info_locked(std::uint32_t slot, const CapturedFormat& cf, std::uint32_t state,
                                std::uint64_t frames)
{
	IpcClient* ipc = g_ipc.load(std::memory_order_acquire);
	if (ipc == nullptr)
	{
		return;
	}
	AudioStreamInfo info{};
	info.is_primary = (slot == 0) ? 1u : 0u;
	info.sample_rate = cf.rate;
	info.channels = static_cast<std::uint16_t>(cf.channels);
	info.bits = static_cast<std::uint16_t>(cf.bits);
	info.format_tag = cf.tag;
	info.frames_rendered = frames;
	info.format_state = state;
	ipc->publish_audio_stream(slot, info);
}

// Publish stream `slot`'s format to its ring, first deciding a guessed sample rate by
// robust measurement (rate_estimator.hpp). Returns true once published (false = no ring
// yet, or a guess still being measured, in which case the caller retries next tick).
// Caller holds g_setup_mutex.
bool publish_stream_format_locked(std::uint32_t slot)
{
	AudioRingHeader* ring = g_rings[slot].load(std::memory_order_acquire);
	if (ring == nullptr || g_stream_formats[slot].rate == 0)
	{
		return false; // no ring attached yet, or no stream in this slot
	}
	if (audio_ring_format_ready(*ring))
	{
		return true; // already published
	}
	CapturedFormat cf = g_stream_formats[slot];
	if (g_stream_rate_guess[slot])
	{
		// Feed this tick's render cadence to the estimator; it only commits on consensus
		// across standard-rate windows, or a low-confidence fallback after enough attempts.
		LARGE_INTEGER now{}, freq{};
		QueryPerformanceCounter(&now);
		QueryPerformanceFrequency(&freq);
		const RateEstimate est = g_rate_estimator[slot].feed(
		    g_streams[slot].frames.load(std::memory_order_relaxed), now.QuadPart, freq.QuadPart);
		if (!est.done)
		{
			return false; // still measuring; caller retries next tick
		}
		const std::uint32_t state = est.confident ? AudioFormat_Measured : AudioFormat_LowConfidence;
		if (est.confident)
		{
			logf("audio stream %u: measured rate %uHz (was guessing %uHz)", slot, est.rate, cf.rate);
		}
		else
		{
			logw("audio stream %u: rate %uHz is a LOW-CONFIDENCE estimate (no consensus) -- verify or "
			     "override",
			     slot, est.rate);
		}
		cf.rate = est.rate;
		g_stream_formats[slot].rate = est.rate; // reflect the decision in the debug/UI snapshot
		g_stream_rate_guess[slot] = false;      // rate decided; channels/bits stay the device assumption
		g_stream_format_state[slot] = state;
		publish_stream_info_locked(slot, cf, state, g_streams[slot].frames.load(std::memory_order_relaxed));
	}
	audio_ring_set_format(*ring, cf.rate, cf.channels, cf.bits, cf.tag, cf.block_align);
	logf("audio stream %u: format %uHz/%uch/%ubit -> ring %p", slot, cf.rate, cf.channels, cf.bits, ring);
	return true;
}

// Apply an operator command (host -> hook via the ring op channel) to stream `slot`:
// re-run the rate measurement, or override the format. Both clear the ring's published
// format so publish_stream_format_locked re-publishes (bumping format_generation, which
// makes the host rebuild its render client at the new format). Caller holds g_setup_mutex.
void apply_audio_op_locked(std::uint32_t slot, const AudioRingOpCmd& cmd)
{
	AudioRingHeader* ring = g_rings[slot].load(std::memory_order_acquire);
	if (ring == nullptr || g_stream_formats[slot].rate == 0)
	{
		return; // no ring / no stream in this slot
	}
	if (cmd.kind == AudioRingOp_Remeasure)
	{
		logw("audio stream %u: operator requested re-measure", slot);
		g_stream_rate_guess[slot] = true;
		g_rate_estimator[slot] = RateEstimator{};
		g_stream_formats[slot].rate = g_mix_format.rate; // back to the device-mix guess while measuring
		g_stream_format_state[slot] = AudioFormat_Measuring;
		g_streams[slot].assumed_format.store(1, std::memory_order_relaxed);
		ring->format_valid.store(0, std::memory_order_release); // force re-publish after measuring
		publish_stream_info_locked(slot, g_stream_formats[slot], AudioFormat_Measuring,
		                           g_streams[slot].frames.load(std::memory_order_relaxed));
	}
	else if (cmd.kind == AudioRingOp_Override)
	{
		CapturedFormat cf;
		cf.rate = cmd.rate;
		cf.channels = cmd.channels;
		cf.bits = cmd.bits;
		cf.tag = cmd.format_tag ? cmd.format_tag : WAVE_FORMAT_PCM;
		cf.block_align = cmd.channels * (cmd.bits / 8);
		if (cf.rate == 0 || cf.channels == 0 || cf.block_align == 0)
		{
			logw("audio stream %u: ignoring invalid override %uHz/%uch/%ubit", slot, cf.rate, cf.channels,
			     cf.bits);
			return;
		}
		logw("audio stream %u: operator override -> %uHz/%uch/%ubit tag=%u", slot, cf.rate, cf.channels,
		     cf.bits, cf.tag);
		g_stream_formats[slot] = cf;
		g_stream_rate_guess[slot] = false;
		g_stream_format_state[slot] = AudioFormat_Override;
		// Keep the over-read clamp on: a too-large operator block is capped to the real
		// buffer (garbled but safe); a correct override makes the clamp a no-op.
		g_streams[slot].assumed_format.store(1, std::memory_order_relaxed);
		g_streams[slot].block_align.store(cf.block_align, std::memory_order_relaxed);
		ring->format_valid.store(0, std::memory_order_release); // re-publish at the new format
		publish_stream_info_locked(slot, cf, AudioFormat_Override,
		                           g_streams[slot].frames.load(std::memory_order_relaxed));
	}
}

// Registers a newly created render client: assigns it a debug slot, marks the
// first as primary (the one we capture), publishes it to HookStatus, and hooks
// the render-client vtable on first sight. `rate_is_guess` is true when `cf` is the
// device mix format assumed for a pre-existing client (its rate is then measured before
// the format is published). Caller holds g_setup_mutex.
void register_render_client_locked(IAudioRenderClient* rc, const CapturedFormat& cf, bool rate_is_guess)
{
	for (std::uint32_t i = 0; i < kMaxAudioStreams; ++i)
	{
		if (g_streams[i].client.load(std::memory_order_relaxed) == rc)
		{
			return; // already tracked
		}
	}

	const std::uint32_t seen = g_streams_seen.fetch_add(1, std::memory_order_relaxed) + 1;
	if (IpcClient* ipc = g_ipc.load(std::memory_order_acquire))
	{
		ipc->set_audio_streams_seen(seen);
	}
	logf("register_render_client: rc=%p seen=%u fmt=%uHz/%uch/%ubit tag=%u block=%u", rc, seen, cf.rate,
	     cf.channels, cf.bits, cf.tag, cf.block_align);

	const std::uint32_t slot = g_registered;
	if (slot >= kMaxAudioStreams)
	{
		return; // more streams than debug slots; counted above, not detailed
	}
	g_registered = slot + 1;

	const std::uint32_t state = rate_is_guess ? AudioFormat_Measuring : AudioFormat_Exact;
	g_stream_formats[slot] = cf;
	g_stream_rate_guess[slot] = rate_is_guess;
	g_stream_format_state[slot] = state;
	g_rate_estimator[slot] = RateEstimator{}; // fresh measurement (used only for a guess)
	g_streams[slot].frames.store(0, std::memory_order_relaxed);
	g_streams[slot].assumed_format.store(rate_is_guess ? 1u : 0u, std::memory_order_relaxed);
	g_streams[slot].block_align.store(cf.block_align, std::memory_order_relaxed); // before client (hot path)
	g_streams[slot].client.store(rc, std::memory_order_release);

	if (rate_is_guess)
	{
		logf("audio stream %u: format unknown (pre-existing client) -> assuming device mix %uHz/%uch/%ubit; "
		     "measuring true rate; channels/bits assumed (verified byte-compatible before capture)",
		     slot, cf.rate, cf.channels, cf.bits);
	}
	else
	{
		logf("audio stream %u: exact format %uHz/%uch/%ubit from the game's Initialize", slot, cf.rate,
		     cf.channels, cf.bits);
	}
	publish_stream_info_locked(slot, cf, state, 0);

	// Publish the format to the stream's ring (if the host has attached one). A guessed
	// rate is measured/corrected inside the helper first, so this may defer until enough
	// audio has rendered to measure -- republish_audio_format retries each worker tick.
	publish_stream_format_locked(slot);
	// GetBuffer/ReleaseBuffer are hooked proactively at install time (the shared
	// vtable covers every render client), so nothing to install per-stream here.
}

// True if `rc` already occupies a tracked debug slot (lock-free scan).
bool stream_tracked(IAudioRenderClient* rc)
{
	for (auto& s : g_streams)
	{
		if (s.client.load(std::memory_order_acquire) == rc)
		{
			return true;
		}
	}
	return false;
}

// Register a render client discovered on the audio thread (we never saw its
// Initialize/GetService — it predates our injection). Uses the device mix format
// as a best guess. Non-blocking: if setup is momentarily busy, retry next call.
void try_register_lazy(IAudioRenderClient* rc)
{
	if (g_have_mix_format.load(std::memory_order_acquire) == 0)
	{
		return;
	}
	std::unique_lock<std::mutex> lock(g_setup_mutex, std::try_to_lock);
	if (!lock.owns_lock())
	{
		return; // another thread is in setup; try again on the next buffer
	}
	if (stream_tracked(rc))
	{
		return; // a concurrent path registered it first
	}
	logf("try_register_lazy: discovered pre-existing render client rc=%p", rc);
	register_render_client_locked(rc, g_mix_format, /*rate_is_guess=*/true);
}

HRESULT STDMETHODCALLTYPE hk_Initialize(IAudioClient* self, AUDCLNT_SHAREMODE mode, DWORD flags,
                                        REFERENCE_TIME buffer_duration, REFERENCE_TIME periodicity,
                                        const WAVEFORMATEX* format, LPCGUID session)
{
	hook_note_call(g_id_initialize);
	const HRESULT hr = g_vh_initialize.original<InitializeFn>()(self, mode, flags, buffer_duration,
	                                                            periodicity, format, session);
	logf("hk_Initialize: client=%p mode=%d flags=0x%lX hr=0x%08lX fmt=%s", self, mode,
	     static_cast<unsigned long>(flags), static_cast<unsigned long>(hr), format ? "yes" : "null");
	if (SUCCEEDED(hr) && format != nullptr)
	{
		std::scoped_lock lock(g_setup_mutex);
		g_client_formats[self] = capture_format(format);
	}
	return hr;
}

HRESULT STDMETHODCALLTYPE hk_GetService(IAudioClient* self, REFIID riid, void** ppv)
{
	hook_note_call(g_id_getservice);
	const HRESULT hr = g_vh_getservice.original<GetServiceFn>()(self, riid, ppv);
	const bool is_render = (riid == __uuidof(IAudioRenderClient));
	logf("hk_GetService: client=%p hr=0x%08lX render_client=%d", self, static_cast<unsigned long>(hr),
	     is_render ? 1 : 0);
	if (SUCCEEDED(hr) && ppv != nullptr && *ppv != nullptr && riid == __uuidof(IAudioRenderClient))
	{
		CapturedFormat cf;
		bool have = false;
		{
			std::scoped_lock lock(g_setup_mutex);
			auto it = g_client_formats.find(self);
			if (it != g_client_formats.end())
			{
				cf = it->second;
				have = true;
			}
		}
		// Fallback for IAudioClient3::InitializeSharedAudioStream (no Initialize
		// format): the shared-mode format is the device mix format.
		if (!have)
		{
			WAVEFORMATEX* mix = nullptr;
			if (SUCCEEDED(self->GetMixFormat(&mix)) && mix != nullptr)
			{
				cf = capture_format(mix);
				have = true;
				CoTaskMemFree(mix);
			}
		}
		if (have)
		{
			std::scoped_lock lock(g_setup_mutex);
			// We saw this client's Initialize (or its shared-mode mix format), so the rate
			// is exact, not a guess.
			register_render_client_locked(static_cast<IAudioRenderClient*>(*ppv), cf, /*rate_is_guess=*/false);
		}
	}
	return hr;
}

void install_audioclient_hooks(IAudioClient* ac)
{
	std::scoped_lock lock(g_setup_mutex);
	if (g_audioclient_hooked)
	{
		return; // shared vtable: hook the first IAudioClient we see, covers all
	}
	g_vh_initialize.install(ac, kIdx_IAudioClient_Initialize, reinterpret_cast<void*>(&hk_Initialize));
	g_vh_getservice.install(ac, kIdx_IAudioClient_GetService, reinterpret_cast<void*>(&hk_GetService));
	g_audioclient_hooked = (static_cast<bool>(g_vh_initialize) && static_cast<bool>(g_vh_getservice));
	logf("install_audioclient_hooks: ac=%p initialize=%d getservice=%d", ac,
	     static_cast<bool>(g_vh_initialize) ? 1 : 0, static_cast<bool>(g_vh_getservice) ? 1 : 0);
}

HRESULT STDMETHODCALLTYPE hk_Activate(IMMDevice* self, REFIID riid, DWORD cls_ctx, PROPVARIANT* params,
                                      void** ppv)
{
	hook_note_call(g_id_activate);
	const HRESULT hr = g_vh_activate.original<ActivateFn>()(self, riid, cls_ctx, params, ppv);
	const bool is_audioclient = (riid == __uuidof(IAudioClient) || riid == __uuidof(IAudioClient2) ||
	                             riid == __uuidof(IAudioClient3));
	logf("hk_Activate: device=%p hr=0x%08lX audioclient=%d", self, static_cast<unsigned long>(hr),
	     is_audioclient ? 1 : 0);
	if (SUCCEEDED(hr) && ppv != nullptr && *ppv != nullptr && is_audioclient)
	{
		install_audioclient_hooks(static_cast<IAudioClient*>(*ppv));
	}
	return hr;
}

} // namespace

namespace
{
// Build the probe COM objects (enumerator -> device -> client -> render) and capture the
// device mix format. Created ONCE and kept for the DLL's lifetime: every instance of a
// coclass shares one vtable, so a toggle then only re-swaps vtable slots on these kept
// objects -- no COM create/destroy churn (which raced AudioSes). Caller holds g_setup_mutex.
bool build_probe_locked()
{
	if (g_self_device != nullptr)
	{
		return true; // already built
	}
	IMMDeviceEnumerator* enumerator = nullptr;
	if (FAILED(CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
	                            __uuidof(IMMDeviceEnumerator), reinterpret_cast<void**>(&enumerator))))
	{
		return false;
	}
	IMMDevice* device = nullptr;
	const HRESULT hr = enumerator->GetDefaultAudioEndpoint(eRender, eConsole, &device);
	enumerator->Release(); // only needed to reach the device
	if (FAILED(hr) || device == nullptr)
	{
		return false;
	}
	g_self_device = device; // kept alive (Activate hook re-installs from its vtable)

	IAudioRenderClient* self_render = nullptr;
	HRESULT ah = device->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr,
	                              reinterpret_cast<void**>(&g_self_client));
	if (SUCCEEDED(ah) && g_self_client != nullptr)
	{
		WAVEFORMATEX* mix = nullptr;
		if (SUCCEEDED(g_self_client->GetMixFormat(&mix)) && mix != nullptr)
		{
			g_mix_format = capture_format(mix);
			g_have_mix_format.store(1, std::memory_order_release);
			constexpr REFERENCE_TIME kBuf = 10 * 10000; // 10 ms; never started
			HRESULT ih = g_self_client->Initialize(AUDCLNT_SHAREMODE_SHARED, 0, kBuf, 0, mix, nullptr);
			if (SUCCEEDED(ih))
			{
				ih = g_self_client->GetService(__uuidof(IAudioRenderClient),
				                               reinterpret_cast<void**>(&self_render));
			}
			logf("build_probe: client init=0x%08lX render=%p mix=%uHz/%uch/%ubit tag=%u",
			     static_cast<unsigned long>(ih), self_render, g_mix_format.rate, g_mix_format.channels,
			     g_mix_format.bits, g_mix_format.tag);
			CoTaskMemFree(mix);
		}
	}
	else
	{
		logf("build_probe: Activate(IAudioClient) failed hr=0x%08lX", static_cast<unsigned long>(ah));
	}
	if (self_render != nullptr)
	{
		g_self_render.store(self_render, std::memory_order_release);
	}
	return true; // device built; render may be null on odd setups (Activate hook still works)
}

// Swap the detours into the shared vtables using the kept probe objects. Caller holds
// g_setup_mutex.
void install_detours_locked()
{
	if (g_self_device == nullptr)
	{
		return;
	}
	g_hook_epoch.fetch_add(1, std::memory_order_release); // new epoch: invalidate any straddling GetBuffer
	g_vh_activate.install(g_self_device, kIdx_IMMDevice_Activate, reinterpret_cast<void*>(&hk_Activate));
	if (IAudioRenderClient* sr = g_self_render.load(std::memory_order_acquire))
	{
		g_vh_initialize.install(g_self_client, kIdx_IAudioClient_Initialize,
		                        reinterpret_cast<void*>(&hk_Initialize));
		g_vh_getservice.install(g_self_client, kIdx_IAudioClient_GetService,
		                        reinterpret_cast<void*>(&hk_GetService));
		g_vh_getbuffer.install(sr, kIdx_IAudioRenderClient_GetBuffer, reinterpret_cast<void*>(&hk_GetBuffer));
		g_vh_releasebuffer.install(sr, kIdx_IAudioRenderClient_ReleaseBuffer,
		                           reinterpret_cast<void*>(&hk_ReleaseBuffer));
		g_audioclient_hooked = (static_cast<bool>(g_vh_initialize) && static_cast<bool>(g_vh_getservice));
	}
	hook_set_installed(g_id_activate, static_cast<bool>(g_vh_activate));
	hook_set_installed(g_id_initialize, static_cast<bool>(g_vh_initialize));
	hook_set_installed(g_id_getservice, static_cast<bool>(g_vh_getservice));
	hook_set_installed(g_id_getbuffer, static_cast<bool>(g_vh_getbuffer));
	hook_set_installed(g_id_releasebuffer, static_cast<bool>(g_vh_releasebuffer));
	logf("install_detours: activate=%d init=%d getsvc=%d getbuf=%d relbuf=%d",
	     static_cast<bool>(g_vh_activate) ? 1 : 0, static_cast<bool>(g_vh_initialize) ? 1 : 0,
	     static_cast<bool>(g_vh_getservice) ? 1 : 0, static_cast<bool>(g_vh_getbuffer) ? 1 : 0,
	     static_cast<bool>(g_vh_releasebuffer) ? 1 : 0);
}
} // namespace

bool install_audio_hooks(IpcClient& ipc, AudioRingHeader* ring)
{
	std::scoped_lock lock(g_setup_mutex);
	g_ipc.store(&ipc, std::memory_order_release);
	g_rings[0].store(ring, std::memory_order_release);
	if (g_vh_activate)
	{
		return true; // detours already installed
	}
	if (g_id_activate < 0) // register the hook-list ids once
	{
		g_id_activate = hook_register("IMMDevice::Activate", HookSubsys_Audio);
		g_id_initialize = hook_register("IAudioClient::Initialize", HookSubsys_Audio);
		g_id_getservice = hook_register("IAudioClient::GetService", HookSubsys_Audio);
		g_id_getbuffer = hook_register("IAudioRenderClient::GetBuffer", HookSubsys_Audio);
		g_id_releasebuffer = hook_register("IAudioRenderClient::ReleaseBuffer", HookSubsys_Audio);
	}
	if (!build_probe_locked()) // builds once; a re-enable reuses the kept probe
	{
		return false;
	}
	install_detours_locked();
	return static_cast<bool>(g_vh_activate);
}

void republish_audio_format()
{
	std::scoped_lock lock(g_setup_mutex);
	for (std::uint32_t i = 0; i < kMaxAudioStreams; ++i)
	{
		AudioRingHeader* ring = g_rings[i].load(std::memory_order_acquire);
		if (ring == nullptr)
		{
			continue;
		}
		// Apply any operator command (re-measure / override) the host posted on this ring.
		AudioRingOpCmd cmd;
		if (audio_ring_poll_op(*ring, g_last_op_seq[i], cmd) != AudioRingOp_None)
		{
			apply_audio_op_locked(i, cmd);
		}
		// Publishes an exact format immediately; a guessed rate is measured first and
		// published once a measurement window completes (retried on the next tick).
		publish_stream_format_locked(i);
	}
}

void set_audio_ring(unsigned index, AudioRingHeader* ring)
{
	if (index >= kMaxAudioStreams)
	{
		return;
	}
	// The worker thread re-attaches every tick (idempotent); only log when the ring
	// pointer actually changes so the log isn't flooded with identical lines.
	AudioRingHeader* const prev = g_rings[index].exchange(ring, std::memory_order_acq_rel);
	if (prev != ring)
	{
		logf("set_audio_ring: index=%u ring=%p capture_enabled=%u", index, ring,
		     ring ? ring->capture_enabled.load(std::memory_order_relaxed) : 0u);
	}
	// The stream may already be registered (game was playing before we injected and
	// before the host created the ring); publish its format so the host stops waiting
	// and consumes the ring instead of falling back to loopback.
	republish_audio_format();
}

void remove_audio_hooks()
{
	std::scoped_lock lock(g_setup_mutex);
	// Restore the vtable slots (so the game's calls go direct again), but KEEP the probe
	// objects alive -- a re-enable just re-swaps the slots, no COM churn. The probe is only
	// released on detach (shutdown_audio_hooks). m_original stays valid (see VtableHook), so
	// an in-flight detour on the audio thread completes safely after the restore.
	g_hook_epoch.fetch_add(1, std::memory_order_release); // new epoch: a later capture won't trust a pre-toggle GetBuffer
	g_vh_releasebuffer.remove();
	g_vh_getbuffer.remove();
	g_vh_getservice.remove();
	g_vh_initialize.remove();
	g_vh_activate.remove();
	g_audioclient_hooked = false;
	hook_set_installed(g_id_activate, false);
	hook_set_installed(g_id_initialize, false);
	hook_set_installed(g_id_getservice, false);
	hook_set_installed(g_id_getbuffer, false);
	hook_set_installed(g_id_releasebuffer, false);

	// The slots are restored above, so no NEW detour will start. Drain any detour still
	// in-flight on the audio thread before clearing the shared state it reads (an initial
	// sleep covers a caller that read the old slot but hasn't entered the guard yet). Bounded
	// so a wedged audio thread can't hang us. Detours are microseconds, so this is ~1-2 ms.
	for (int spins = 0; spins < 200; ++spins)
	{
		Sleep(1);
		if (g_detours_active.load(std::memory_order_acquire) == 0)
		{
			break;
		}
	}

	// Clear per-stream tracking (the game's streams re-register lazily on a re-enable).
	g_registered = 0;
	g_streams_seen.store(0, std::memory_order_relaxed);
	g_frames_captured.store(0, std::memory_order_relaxed);
	g_frames_silenced.store(0, std::memory_order_relaxed);
	for (std::uint32_t i = 0; i < kMaxAudioStreams; ++i)
	{
		g_streams[i].client.store(nullptr, std::memory_order_relaxed);
		g_streams[i].frames.store(0, std::memory_order_relaxed);
		g_streams[i].block_align.store(0, std::memory_order_relaxed);
		g_streams[i].assumed_format.store(0, std::memory_order_relaxed);
		g_stream_formats[i] = CapturedFormat{};
		g_stream_rate_guess[i] = false;
		g_stream_format_state[i] = AudioFormat_Unknown;
		g_rate_estimator[i] = RateEstimator{};
		g_last_op_seq[i] = 0;
		g_rings[i].store(nullptr, std::memory_order_release);
	}
	g_client_formats.clear();
	g_ipc.store(nullptr, std::memory_order_release);
}

void shutdown_audio_hooks()
{
	remove_audio_hooks(); // restore vtables + clear state (takes the lock)
	// Now safe to release the kept probe objects (called only on DLL detach).
	std::scoped_lock lock(g_setup_mutex);
	if (IAudioRenderClient* sr = g_self_render.exchange(nullptr, std::memory_order_acq_rel))
	{
		sr->Release();
	}
	if (g_self_client != nullptr)
	{
		g_self_client->Release();
		g_self_client = nullptr;
	}
	if (g_self_device != nullptr)
	{
		g_self_device->Release();
		g_self_device = nullptr;
	}
	g_have_mix_format.store(0, std::memory_order_relaxed);
}

std::uint64_t audio_frames_silenced()
{
	return g_frames_silenced.load(std::memory_order_relaxed);
}

std::uint64_t audio_frames_captured()
{
	return g_frames_captured.load(std::memory_order_relaxed);
}

std::uint32_t audio_streams_seen()
{
	return g_streams_seen.load(std::memory_order_relaxed);
}

} // namespace coop::hook