// Configurable WASAPI sine-tone render source, shared by coop_tone.exe and the audio
// render-hook self-test. Opens a shared-mode render client at a requested format
// (sample rate / channels / bits / float vs PCM) using AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM,
// so it can render formats that differ from the device mix format -- exactly the case a
// game rendering 44100 Hz on a 48000 Hz endpoint creates, which the hook must detect.
#pragma once

#include <cmath>

#include <windows.h>

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

namespace coop::tone {

inline constexpr double kTwoPi = 6.283185307179586;

// A field left 0 resolves to the device mix format's value (so {} = play at the device
// format). `is_float` only applies when `bits` is set (16 -> PCM, 32 -> float by default).
struct ToneFormat {
	unsigned rate = 0;
	unsigned channels = 0;
	unsigned bits = 0;
	bool is_float = false;
};

class ToneSource {
  public:
	~ToneSource() { close(); }

	// Open + start a render client at `want` (0 fields resolve to the device mix format,
	// AUTOCONVERTPCM lets a non-device format be rendered). Returns false if the endpoint
	// or that specific format isn't available (the caller treats that as a per-format skip).
	bool open(const ToneFormat& want, double freq_hz = 440.0)
	{
		if (FAILED(CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL, __uuidof(IMMDeviceEnumerator),
		                            reinterpret_cast<void**>(&enum_)))) {
			return false;
		}
		if (FAILED(enum_->GetDefaultAudioEndpoint(eRender, eConsole, &endpoint_))) {
			return false;
		}
		if (FAILED(
		        endpoint_->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, reinterpret_cast<void**>(&client_)))) {
			return false;
		}

		WAVEFORMATEX* mix = nullptr;
		if (FAILED(client_->GetMixFormat(&mix)) || mix == nullptr) {
			return false;
		}
		resolve_format(want, mix);
		CoTaskMemFree(mix);

		WAVEFORMATEXTENSIBLE wfx{};
		build_waveformat(wfx);
		auto* fmt = reinterpret_cast<WAVEFORMATEX*>(&wfx);

		event_ = CreateEventW(nullptr, FALSE, FALSE, nullptr);
		constexpr REFERENCE_TIME kBuffer = 30 * 10000; // 30 ms
		// AUTOCONVERTPCM makes a shared-mode client render a non-device format (the audio
		// engine resamples to the endpoint), exactly like the games that need rate detection.
		const DWORD flags = AUDCLNT_STREAMFLAGS_EVENTCALLBACK | AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM
		                    | AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY;
		if (FAILED(client_->Initialize(AUDCLNT_SHAREMODE_SHARED, flags, kBuffer, 0, fmt, nullptr))) {
			return false;
		}
		client_->SetEventHandle(event_);
		if (FAILED(client_->GetService(__uuidof(IAudioRenderClient), reinterpret_cast<void**>(&render_)))) {
			return false;
		}
		client_->GetBufferSize(&buffer_frames_);

		step_ = kTwoPi * freq_hz / static_cast<double>(fmt_.rate);
		// Optional: give each channel genuinely different content (a per-channel frequency scale),
		// so a downstream test can *recover* the channel count by correlation (identical channels
		// are ambiguous: 2ch@R looks like 1ch@2R). Off by default -> the usual single-tone source.
		char d[2] = {};
		if (GetEnvironmentVariableA("COOP_TONE_DISTINCT_CH", d, sizeof(d)) > 0 && d[0] == '1') {
			distinct_ = true;
		}
		write(buffer_frames_); // pre-roll
		client_->Start();
		return true;
	}

	// Wait up to `timeout_ms` for the buffer event, then refill. Returns false on a
	// timeout/error (the caller keeps looping on its own wall clock).
	bool render_step(DWORD timeout_ms)
	{
		if (render_ == nullptr) {
			return false;
		}
		if (WaitForSingleObject(event_, timeout_ms) != WAIT_OBJECT_0) {
			return false;
		}
		UINT32 padding = 0;
		if (FAILED(client_->GetCurrentPadding(&padding))) {
			return false;
		}
		write(buffer_frames_ - padding);
		return true;
	}

	const ToneFormat& format() const { return fmt_; }
	bool is_open() const { return render_ != nullptr; }

	void close()
	{
		if (client_) {
			client_->Stop();
		}
		rel(render_);
		rel(client_);
		rel(endpoint_);
		rel(enum_);
		if (event_) {
			CloseHandle(event_);
			event_ = nullptr;
		}
	}

  private:
	template <typename T>
	static void rel(T*& p)
	{
		if (p) {
			p->Release();
			p = nullptr;
		}
	}

	void resolve_format(const ToneFormat& want, const WAVEFORMATEX* mix)
	{
		fmt_.rate = want.rate ? want.rate : mix->nSamplesPerSec;
		fmt_.channels = want.channels ? want.channels : mix->nChannels;
		if (want.bits) {
			fmt_.bits = want.bits;
			fmt_.is_float = want.is_float;
		} else {
			fmt_.bits = mix->wBitsPerSample;
			fmt_.is_float = mix->wFormatTag == WAVE_FORMAT_IEEE_FLOAT
			                || (mix->wFormatTag == WAVE_FORMAT_EXTENSIBLE
			                    && reinterpret_cast<const WAVEFORMATEXTENSIBLE*>(mix)->SubFormat
			                           == KSDATAFORMAT_SUBTYPE_IEEE_FLOAT);
		}
		float_ = fmt_.is_float;
	}

	void build_waveformat(WAVEFORMATEXTENSIBLE& wfx)
	{
		const WORD block = static_cast<WORD>(fmt_.channels * (fmt_.bits / 8));
		wfx.Format.nChannels = static_cast<WORD>(fmt_.channels);
		wfx.Format.nSamplesPerSec = fmt_.rate;
		wfx.Format.wBitsPerSample = static_cast<WORD>(fmt_.bits);
		wfx.Format.nBlockAlign = block;
		wfx.Format.nAvgBytesPerSec = block * fmt_.rate;
		if (fmt_.channels > 2 || fmt_.bits > 16) {
			wfx.Format.wFormatTag = WAVE_FORMAT_EXTENSIBLE;
			wfx.Format.cbSize = sizeof(WAVEFORMATEXTENSIBLE) - sizeof(WAVEFORMATEX);
			wfx.Samples.wValidBitsPerSample = static_cast<WORD>(fmt_.bits);
			switch (fmt_.channels) {
			case 6:
				wfx.dwChannelMask = 0x3F;
				break;
			case 8:
				wfx.dwChannelMask = 0xFF;
				break;
			default:
				wfx.dwChannelMask = (1u << fmt_.channels) - 1u;
				break;
			}
			wfx.SubFormat = float_ ? KSDATAFORMAT_SUBTYPE_IEEE_FLOAT : KSDATAFORMAT_SUBTYPE_PCM;
		} else {
			wfx.Format.wFormatTag = float_ ? WAVE_FORMAT_IEEE_FLOAT : WAVE_FORMAT_PCM;
			wfx.Format.cbSize = 0;
		}
	}

	void write(UINT32 frames)
	{
		BYTE* data = nullptr;
		if (frames == 0 || render_ == nullptr || FAILED(render_->GetBuffer(frames, &data))) {
			return;
		}
		for (UINT32 i = 0; i < frames; ++i) {
			const double s = std::sin(phase_) * 0.25; // -12 dB, gentle
			phase_ += step_;
			if (phase_ > kTwoPi) {
				phase_ -= kTwoPi;
			}
			for (unsigned c = 0; c < fmt_.channels; ++c) {
				double sc = s;
				if (distinct_ && c < 8) {
					// Each channel at its own frequency scale -> genuinely different content.
					sc = std::sin(phase_c_[c]) * 0.25;
					phase_c_[c] += step_ * (1.0 + 0.37 * static_cast<double>(c));
					if (phase_c_[c] > kTwoPi) {
						phase_c_[c] -= kTwoPi;
					}
				}
				if (float_) {
					reinterpret_cast<float*>(data)[i * fmt_.channels + c] = static_cast<float>(sc);
				} else {
					reinterpret_cast<INT16*>(data)[i * fmt_.channels + c] = static_cast<INT16>(sc * 32767.0);
				}
			}
		}
		render_->ReleaseBuffer(frames, 0);
	}

	IMMDeviceEnumerator* enum_ = nullptr;
	IMMDevice* endpoint_ = nullptr;
	IAudioClient* client_ = nullptr;
	IAudioRenderClient* render_ = nullptr;
	HANDLE event_ = nullptr;
	UINT32 buffer_frames_ = 0;
	ToneFormat fmt_;
	bool float_ = false;
	double phase_ = 0.0;
	double step_ = 0.0;
	bool distinct_ = false;  // per-channel distinct content (recoverable channel count)
	double phase_c_[8] = {}; // per-channel phase when distinct_
};

} // namespace coop::tone