// 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). audio_correlate_layout.hpp (step b) reuses these 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;
}

} // namespace coop