// Recover a pre-existing render stream's true audio format by *correlating* the two capture
// paths, instead of guessing.
//
// When we attach to an already-running game we never saw its IAudioClient::Initialize, so the
// render-hook assumes the device mix format and measures only the sample rate from the render
// cadence -- which can be wrong on a jittery game (intermittent pitch shift). But during the
// measurement window the game is still audible, so we have BOTH signals of the same audio:
//   * the render-hook capture  -- pre-mix, at the *unknown* format,
//   * the process-loopback capture -- post-mix, at the *known* device format.
// The loopback is just the hook signal resampled by WASAPI's AUTOCONVERTPCM from the stream's
// true rate to the device rate. So if we resample the hook stream by a candidate rate up to the
// device rate and it lines up with the loopback over the whole window (no drift), that candidate
// is the truth. A wrong rate time-warps the hook stream, so a single alignment can't hold across
// the window and the correlation collapses.
//
// This header is the pure, headless-testable core (no devices, no WASAPI). The host downmixes the
// two captures to mono float, calls correlate_rate(), and feeds the result into the existing rate
// path (publish / override). correlate_format() below builds on the same helpers to also recover
// channels + bit depth by trying candidate de-interleavings.
#pragma once

#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>

namespace coop
{

// The standard sample rates a shared-mode WASAPI stream realistically uses. Candidates are this
// set; a non-standard true rate is out of scope (and would show as low-confidence either way).
inline const std::vector<unsigned>& standard_audio_rates()
{
	static const std::vector<unsigned> rates = {32000, 44100, 48000, 88200, 96000};
	return rates;
}

struct RateCorrelation
{
	bool ok = false;       // a confident pick was made (winner clears the threshold AND beats the runner-up)
	unsigned rate = 0;     // best candidate rate (Hz)
	double score = 0.0;    // alignment score of the winner, in [0,1] (1 = perfect)
	double runner_up = 0.0; // score of the second-best candidate (for separation)
};

namespace correlate_detail
{

// Average interleaved float frames down to a single mono channel.
inline void downmix(const float* interleaved, std::size_t frames, unsigned channels, std::vector<float>& out)
{
	out.resize(frames);
	if (channels == 0)
	{
		channels = 1;
	}
	for (std::size_t i = 0; i < frames; ++i)
	{
		float sum = 0.0f;
		for (unsigned c = 0; c < channels; ++c)
		{
			sum += interleaved[i * channels + c];
		}
		out[i] = sum / static_cast<float>(channels);
	}
}

// Linear-resample a mono signal from src_rate to dst_rate.
inline void resample_linear(const std::vector<float>& in, unsigned src_rate, unsigned dst_rate,
                            std::vector<float>& out)
{
	if (src_rate == 0 || dst_rate == 0 || in.empty())
	{
		out.clear();
		return;
	}
	if (src_rate == dst_rate)
	{
		out = in;
		return;
	}
	const double step = static_cast<double>(src_rate) / static_cast<double>(dst_rate);
	const std::size_t out_n = static_cast<std::size_t>(static_cast<double>(in.size()) / step);
	out.resize(out_n);
	for (std::size_t i = 0; i < out_n; ++i)
	{
		const double pos = static_cast<double>(i) * step;
		const std::size_t j = static_cast<std::size_t>(pos);
		const double frac = pos - static_cast<double>(j);
		const float a = in[j];
		const float b = (j + 1 < in.size()) ? in[j + 1] : a;
		out[i] = a + static_cast<float>(frac) * (b - a);
	}
}

// Box-decimate a mono signal from `rate` down to ~corr_rate for a cheap, content-preserving
// alignment search (the envelope/content alignment doesn't need full bandwidth).
inline void decimate(const std::vector<float>& in, unsigned rate, unsigned corr_rate, std::vector<float>& out)
{
	if (rate <= corr_rate || in.empty())
	{
		out = in;
		return;
	}
	const double factor = static_cast<double>(rate) / static_cast<double>(corr_rate);
	const std::size_t out_n = static_cast<std::size_t>(static_cast<double>(in.size()) / factor);
	out.resize(out_n);
	for (std::size_t i = 0; i < out_n; ++i)
	{
		const std::size_t lo = static_cast<std::size_t>(static_cast<double>(i) * factor);
		std::size_t hi = static_cast<std::size_t>(static_cast<double>(i + 1) * factor);
		if (hi <= lo)
		{
			hi = lo + 1;
		}
		if (hi > in.size())
		{
			hi = in.size();
		}
		float sum = 0.0f;
		for (std::size_t k = lo; k < hi; ++k)
		{
			sum += in[k];
		}
		out[i] = sum / static_cast<float>(hi - lo);
	}
}

// Zero-mean, unit-norm cross-correlation of a vs b over [start, start+len), with b shifted by lag.
// Returns a value in [-1, 1]; out-of-range samples are skipped. ~0 when the two don't align.
inline double ncc(const std::vector<float>& a, const std::vector<float>& b, long lag, std::size_t start,
                  std::size_t len)
{
	double sa = 0.0, sb = 0.0;
	std::size_t n = 0;
	for (std::size_t i = start; i < start + len && i < a.size(); ++i)
	{
		const long bi = static_cast<long>(i) + lag;
		if (bi < 0 || static_cast<std::size_t>(bi) >= b.size())
		{
			continue;
		}
		sa += a[i];
		sb += b[bi];
		++n;
	}
	if (n < 8)
	{
		return 0.0;
	}
	const double ma = sa / static_cast<double>(n);
	const double mb = sb / static_cast<double>(n);
	double num = 0.0, da = 0.0, db = 0.0;
	for (std::size_t i = start; i < start + len && i < a.size(); ++i)
	{
		const long bi = static_cast<long>(i) + lag;
		if (bi < 0 || static_cast<std::size_t>(bi) >= b.size())
		{
			continue;
		}
		const double xa = a[i] - ma;
		const double xb = b[bi] - mb;
		num += xa * xb;
		da += xa * xa;
		db += xb * xb;
	}
	if (da < 1e-9 || db < 1e-9)
	{
		return 0.0;
	}
	return num / std::sqrt(da * db);
}

// Alignment score of two same-rate mono signals: find the single best lag over the whole window,
// then require that lag to hold in BOTH an early and a late segment (drift detection). The score
// is the weaker of the two segment correlations, so a rate that only lines up at the start (a
// wrong rate, which time-warps and drifts) scores low while the true rate scores high.
inline double aligned_score(const std::vector<float>& a, const std::vector<float>& b, unsigned rate)
{
	const std::size_t n = a.size() < b.size() ? a.size() : b.size();
	if (n < rate / 10) // need at least ~100 ms of overlap to judge
	{
		return 0.0;
	}
	const long max_lag = static_cast<long>(rate / 8); // search +/-125 ms of capture-path latency skew
	// Coarse global lag from the middle half of the window.
	const std::size_t mid_start = n / 4;
	const std::size_t mid_len = n / 2;
	double best = -2.0;
	long best_lag = 0;
	for (long lag = -max_lag; lag <= max_lag; ++lag)
	{
		const double c = ncc(a, b, lag, mid_start, mid_len);
		if (c > best)
		{
			best = c;
			best_lag = lag;
		}
	}
	// Re-evaluate that lag in an early and a late third: the true rate holds; a drifting (wrong)
	// rate does not.
	const std::size_t third = n / 3;
	const double early = ncc(a, b, best_lag, 0, third);
	const double late = ncc(a, b, best_lag, 2 * third, third);
	const double weaker = early < late ? early : late;
	return weaker < 0.0 ? 0.0 : weaker;
}

} // namespace correlate_detail

// Determine the hook stream's true sample rate by resampling it by each candidate rate up to the
// known device (loopback) rate and scoring how well it aligns with the loopback across the window.
// hook_mono / loop_mono are mono float (caller downmixes). `min_score` is the absolute alignment
// floor and `separation` the ratio by which the winner must beat the runner-up to be `ok`.
inline RateCorrelation correlate_rate(const std::vector<float>& hook_mono, const std::vector<float>& loop_mono,
                                      unsigned device_rate, const std::vector<unsigned>& candidates,
                                      double min_score = 0.55, double separation = 1.2)
{
	using namespace correlate_detail;
	RateCorrelation result;
	if (hook_mono.empty() || loop_mono.empty() || device_rate == 0)
	{
		return result;
	}
	constexpr unsigned kCorrRate = 8000; // alignment search rate (Nyquist 4 kHz -- plenty for content)

	std::vector<float> loop_ds;
	decimate(loop_mono, device_rate, kCorrRate, loop_ds);

	double best = -1.0, second = -1.0;
	unsigned best_rate = 0;
	std::vector<float> resampled, hook_ds;
	for (unsigned cand : candidates)
	{
		resample_linear(hook_mono, cand, device_rate, resampled); // treat hook as sampled at `cand`
		decimate(resampled, device_rate, kCorrRate, hook_ds);
		const double s = aligned_score(hook_ds, loop_ds, kCorrRate);
		if (s > best)
		{
			second = best;
			best = s;
			best_rate = cand;
		}
		else if (s > second)
		{
			second = s;
		}
	}

	result.rate = best_rate;
	result.score = best < 0.0 ? 0.0 : best;
	result.runner_up = second < 0.0 ? 0.0 : second;
	result.ok = result.score >= min_score && (result.runner_up <= 1e-6 || result.score >= result.runner_up * separation);
	return result;
}

// --- Channels + bit-depth recovery -------------------------------------------------------------
//
// The rate step assumes the hook bytes are de-interleaved at the device channel/bit layout. When a
// game renders a DIFFERENT layout than the device (e.g. stereo float on a 7.1 endpoint, or 16-bit
// PCM), that assumption garbles the waveform and the rate won't lock. We can't measure the layout
// (AUTOCONVERTPCM hides the stride), but we can RECOVER it: interpret the raw hook bytes under each
// candidate layout, run the rate correlation, and keep whichever (layout, rate) aligns with the
// loopback -- a wrong de-interleaving is noise and won't correlate.

inline constexpr unsigned kWaveFormatPcm = 1;   // WAVE_FORMAT_PCM
inline constexpr unsigned kWaveFormatFloat = 3; // WAVE_FORMAT_IEEE_FLOAT

struct LayoutCandidate
{
	unsigned channels;
	unsigned bits;
	unsigned tag; // kWaveFormatPcm / kWaveFormatFloat
};

// Candidate de-interleavings a shared-mode WASAPI render stream realistically uses: float32 and
// 16-bit PCM, across the common channel counts. Ordered most-likely-first.
inline const std::vector<LayoutCandidate>& standard_audio_layouts()
{
	static const std::vector<LayoutCandidate> v = {
	    {2, 32, kWaveFormatFloat}, {1, 32, kWaveFormatFloat}, {6, 32, kWaveFormatFloat},
	    {8, 32, kWaveFormatFloat}, {4, 32, kWaveFormatFloat}, {2, 16, kWaveFormatPcm},
	    {1, 16, kWaveFormatPcm},   {6, 16, kWaveFormatPcm},   {8, 16, kWaveFormatPcm},
	    {4, 16, kWaveFormatPcm},
	};
	return v;
}

struct FormatCorrelation
{
	bool ok = false;
	unsigned rate = 0;
	unsigned channels = 0;
	unsigned bits = 0;
	unsigned tag = 0;
	double score = 0.0;
	double runner_up = 0.0;
};

// The hook can't know a guessed stream's real frame size, so its verify tap pushes each render
// buffer padded to the device block (`stride`) and prefixes it with the real frame `count`. That
// padding is stale staging-buffer bytes, so the host must extract the real `count*real_block` bytes
// per buffer (and concatenate) before de-interleaving -- otherwise the padding scrambles the audio.
// This carries that self-describing capture: `bytes` holds counts[i]*stride bytes per chunk.
struct ChunkedCapture
{
	unsigned stride = 0;               // bytes per frame as pushed (the guessed/device block_align)
	std::vector<std::uint32_t> counts; // real frame count of each chunk
	std::vector<std::uint8_t> bytes;   // concatenated, counts[i]*stride bytes per chunk
};

namespace correlate_detail
{
// De-interleave raw bytes under (channels/bits/tag) and average to mono float.
inline void decode_layout(const std::uint8_t* bytes, std::size_t n, const LayoutCandidate& fmt,
                          std::vector<float>& mono)
{
	mono.clear();
	const unsigned ch = fmt.channels == 0 ? 1 : fmt.channels;
	const unsigned bps = fmt.bits / 8;
	if (bps == 0)
	{
		return;
	}
	const std::size_t frame = static_cast<std::size_t>(ch) * bps;
	const std::size_t frames = n / frame;
	mono.resize(frames);
	const bool is_float = fmt.tag == kWaveFormatFloat;
	for (std::size_t i = 0; i < frames; ++i)
	{
		double sum = 0.0;
		for (unsigned c = 0; c < ch; ++c)
		{
			const std::uint8_t* p = bytes + i * frame + static_cast<std::size_t>(c) * bps;
			float s = 0.0f;
			if (is_float && fmt.bits == 32)
			{
				std::memcpy(&s, p, 4);
			}
			else if (fmt.bits == 16)
			{
				std::int16_t v;
				std::memcpy(&v, p, 2);
				s = v / 32768.0f;
			}
			else if (fmt.bits == 32)
			{
				std::int32_t v;
				std::memcpy(&v, p, 4);
				s = static_cast<float>(v / 2147483648.0);
			}
			sum += s;
		}
		mono[i] = static_cast<float>(sum / ch);
	}
}
} // namespace correlate_detail

// Recover BOTH the layout and the rate of a guessed stream from its (padded, self-describing) hook
// capture: for each candidate layout, extract the real count*real_block bytes from each padded
// chunk, de-interleave to mono, and run the rate correlation against the known-format loopback,
// keeping the (layout, rate) that aligns best. `ok` when the winner clears the alignment floor and
// clearly beats the runner-up (so a coincidental partial match is rejected).
//
// `min_margin` is an ABSOLUTE gap (not a ratio): the true layout scores near-perfectly while a
// truly-ambiguous alternative (e.g. 1ch@2R vs 2ch@R when the channels carry the same content)
// scores within ~0.001, so requiring the winner to clear the runner-up by a fixed margin cleanly
// separates "recovered" from "genuinely ambiguous, don't guess".
inline FormatCorrelation correlate_format(const ChunkedCapture& hook, const std::vector<float>& loop_mono,
                                          unsigned device_rate, const std::vector<unsigned>& rates,
                                          const std::vector<LayoutCandidate>& layouts, double min_score = 0.55,
                                          double min_margin = 0.04)
{
	FormatCorrelation result;
	if (hook.stride == 0 || hook.counts.empty() || loop_mono.empty() || device_rate == 0)
	{
		return result;
	}
	double best = -1.0, second = -1.0;
	std::vector<std::uint8_t> clean;
	std::vector<float> hook_mono;
	for (const LayoutCandidate& layout : layouts)
	{
		const unsigned real_block = layout.channels * (layout.bits / 8);
		if (real_block == 0 || real_block > hook.stride)
		{
			continue; // can't extract a frame larger than what was pushed (the guess is the max)
		}
		// Pull the real count*real_block bytes out of each padded chunk and concatenate -> contiguous
		// audio for this candidate layout (the padding, which is stale staging bytes, is dropped).
		clean.clear();
		std::size_t off = 0;
		for (std::uint32_t count : hook.counts)
		{
			const std::size_t chunk_bytes = static_cast<std::size_t>(count) * hook.stride;
			const std::size_t take = static_cast<std::size_t>(count) * real_block;
			if (off + chunk_bytes <= hook.bytes.size())
			{
				clean.insert(clean.end(), hook.bytes.begin() + off, hook.bytes.begin() + off + take);
			}
			off += chunk_bytes;
		}
		correlate_detail::decode_layout(clean.data(), clean.size(), layout, hook_mono);
		if (hook_mono.size() < device_rate / 5)
		{
			continue; // this layout yields too little audio to judge
		}
		const RateCorrelation rc = correlate_rate(hook_mono, loop_mono, device_rate, rates, /*min_score=*/0.0,
		                                           /*separation=*/1.0);
		if (rc.score > best)
		{
			second = best;
			best = rc.score;
			result.rate = rc.rate;
			result.channels = layout.channels;
			result.bits = layout.bits;
			result.tag = layout.tag;
		}
		else if (rc.score > second)
		{
			second = rc.score;
		}
	}
	result.score = best < 0.0 ? 0.0 : best;
	result.runner_up = second < 0.0 ? 0.0 : second;
	result.ok = result.score >= min_score && (result.runner_up <= 1e-6 || result.score - result.runner_up >= min_margin);
	return result;
}

} // namespace coop