#include "audio/audio_format_verifier.hpp"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>

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

#include "audio/process_loopback_capture.hpp"
#include "coop/audio_correlate.hpp"

namespace coop
{
namespace
{

// Resolve a (possibly EXTENSIBLE) WAVEFORMATEX to scalar channels / bits / tag.
struct ScalarFormat
{
	unsigned rate = 0;
	unsigned channels = 0;
	unsigned bits = 0;
	unsigned tag = 0; // WAVE_FORMAT_PCM / _IEEE_FLOAT
};

ScalarFormat resolve(const WAVEFORMATEX* wfx)
{
	ScalarFormat f;
	f.rate = wfx->nSamplesPerSec;
	f.channels = wfx->nChannels;
	f.bits = wfx->wBitsPerSample;
	f.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)
		{
			f.tag = WAVE_FORMAT_IEEE_FLOAT;
		}
		else if (ext->SubFormat == KSDATAFORMAT_SUBTYPE_PCM)
		{
			f.tag = WAVE_FORMAT_PCM;
		}
	}
	return f;
}

// Decode interleaved PCM (`channels`/`bits`/`tag`) into per-channel float samples, then average to
// mono. Handles float32 and 16/32-bit PCM (the formats WASAPI shared-mode streams use).
std::vector<float> to_mono(const std::vector<BYTE>& bytes, const ScalarFormat& fmt)
{
	std::vector<float> mono;
	const unsigned ch = fmt.channels == 0 ? 1 : fmt.channels;
	const unsigned bps = fmt.bits / 8;
	if (bps == 0)
	{
		return mono;
	}
	const std::size_t frame = static_cast<std::size_t>(ch) * bps;
	const std::size_t frames = bytes.size() / frame;
	mono.resize(frames);
	const bool is_float = fmt.tag == WAVE_FORMAT_IEEE_FLOAT;
	for (std::size_t i = 0; i < frames; ++i)
	{
		double sum = 0.0;
		for (unsigned c = 0; c < ch; ++c)
		{
			const BYTE* p = bytes.data() + 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);
	}
	return mono;
}

void drain_ring(AudioRingHeader& ring, std::vector<BYTE>& scratch)
{
	while (audio_ring_pop(ring, scratch.data(), static_cast<std::uint32_t>(scratch.size())) > 0)
	{
	}
}

// Parse the hook's self-describing verify stream -- a sequence of [u32 frame-count][count*stride
// bytes] chunks -- into a ChunkedCapture. `stride` is the device block_align (what the hook padded
// each buffer to). Stops at the first truncated/garbled record.
ChunkedCapture parse_chunks(const std::vector<BYTE>& raw, unsigned stride)
{
	ChunkedCapture cap;
	cap.stride = stride;
	if (stride == 0)
	{
		return cap;
	}
	std::size_t off = 0;
	while (off + sizeof(std::uint32_t) <= raw.size())
	{
		std::uint32_t count = 0;
		std::memcpy(&count, raw.data() + off, sizeof(count));
		off += sizeof(count);
		const std::size_t payload = static_cast<std::size_t>(count) * stride;
		if (count == 0 || off + payload > raw.size())
		{
			break; // truncated or garbled -> stop
		}
		cap.counts.push_back(count);
		cap.bytes.insert(cap.bytes.end(), raw.data() + off, raw.data() + off + payload);
		off += payload;
	}
	return cap;
}

} // namespace

FormatVerification verify_stream_format(DWORD pid, AudioRingHeader* ring, unsigned window_ms, bool recover_layout)
{
	(void)recover_layout; // step (b) extends this; step (a) recovers the rate only
	FormatVerification result;
	if (ring == nullptr)
	{
		return result;
	}

	WAVEFORMATEX* dev_wfx = default_render_format();
	if (dev_wfx == nullptr)
	{
		return result;
	}
	const ScalarFormat dev = resolve(dev_wfx);

	// Ask the hook to push the guessed stream's pre-mix bytes (no silence) while it measures, and
	// clear any stale ring contents so we only collect this window.
	std::vector<BYTE> scratch(64 * 1024);
	drain_ring(*ring, scratch);
	ring->verify_capture.store(1, std::memory_order_release);

	// Capture the post-mix loopback in parallel (this is the known-format ground truth).
	std::vector<BYTE> loop_bytes;
	ProcessLoopbackCapture loop;
	const std::uint32_t loop_block = dev_wfx->nBlockAlign;
	loop.start(pid, dev_wfx, [&](const BYTE* data, std::uint32_t frames, bool silent) {
		if (!silent && data != nullptr)
		{
			loop_bytes.insert(loop_bytes.end(), data, data + static_cast<std::size_t>(frames) * loop_block);
		}
	});

	// Pull the hook's pre-mix bytes out of the ring across the window.
	std::vector<BYTE> hook_bytes;
	const DWORD end = GetTickCount() + window_ms;
	while (GetTickCount() < end)
	{
		std::uint32_t n = 0;
		while ((n = audio_ring_pop(*ring, scratch.data(), static_cast<std::uint32_t>(scratch.size()))) > 0)
		{
			hook_bytes.insert(hook_bytes.end(), scratch.data(), scratch.data() + n);
		}
		Sleep(10);
	}
	std::uint32_t n = 0;
	while ((n = audio_ring_pop(*ring, scratch.data(), static_cast<std::uint32_t>(scratch.size()))) > 0)
	{
		hook_bytes.insert(hook_bytes.end(), scratch.data(), scratch.data() + n);
	}

	loop.stop();
	ring->verify_capture.store(0, std::memory_order_release);
	drain_ring(*ring, scratch); // leave the ring clean for the real capture that follows

	// The loopback is at the KNOWN device layout. Decode it to mono ground truth. The hook stream
	// is self-describing chunks ([count][padded payload]); parse them at the device block.
	const std::vector<float> loop_mono = to_mono(loop_bytes, dev);
	const unsigned dev_block = dev_wfx->nBlockAlign;
	const ChunkedCapture cap = parse_chunks(hook_bytes, dev_block);
	const std::size_t need = dev.rate / 5; // require >= ~200 ms of usable audio on both sides
	const std::size_t hook_frames = dev_block != 0 ? cap.bytes.size() / dev_block : 0;
	CoTaskMemFree(dev_wfx);

	if (const char* dbg = std::getenv("COOP_VERIFY_DEBUG"); dbg != nullptr && dbg[0] == '1')
	{
		std::fprintf(stderr, "[verify] dev=%uHz/%uch/%ubit blk=%u chunks=%zu hook_frames=%zu loop=%zu layout=%d\n",
		             dev.rate, dev.channels, dev.bits, dev_block, cap.counts.size(), hook_frames, loop_mono.size(),
		             recover_layout ? 1 : 0);
	}
	if (loop_mono.size() < need || hook_frames < need)
	{
		return result; // not enough non-silent audio captured (game quiet, or stream wasn't a guess)
	}

	if (recover_layout)
	{
		// Step (b): recover channels + bit depth too, by trying candidate de-interleavings of the
		// (de-padded) hook bytes and keeping whichever (layout, rate) correlates with the loopback.
		const FormatCorrelation fc =
		    correlate_format(cap, loop_mono, dev.rate, standard_audio_rates(), standard_audio_layouts());
		result.ok = fc.ok;
		result.rate = fc.rate;
		result.score = fc.score;
		result.layout_ok = fc.ok;
		result.channels = fc.channels;
		result.bits = fc.bits;
		result.format_tag = fc.tag;
	}
	else
	{
		// Step (a): rate only, assuming the hook layout matches the device (common stereo case), so
		// the de-padded payload is already clean device-layout audio.
		const std::vector<float> hook_mono = to_mono(cap.bytes, dev);
		const RateCorrelation rc = correlate_rate(hook_mono, loop_mono, dev.rate, standard_audio_rates());
		result.ok = rc.ok;
		result.rate = rc.rate;
		result.score = rc.score;
		result.channels = dev.channels;
		result.bits = dev.bits;
		result.format_tag = dev.tag;
	}
	return result;
}

} // namespace coop