#include "audio/audio_loopback.hpp"

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>

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

#include "audio/audio_mix.hpp"
#include "audio/process_loopback_capture.hpp"

namespace coop
{
namespace
{

// Single-producer/single-consumer byte FIFO guarded by a mutex (the capture
// thread pushes, the render thread pops). Overflow drops the oldest samples.
struct ByteRing
{
	std::mutex mutex;
	std::vector<BYTE> buf;
	size_t head = 0;
	size_t count = 0;

	void init(size_t capacity)
	{
		buf.assign(capacity, 0);
		head = 0;
		count = 0;
	}

	void drop_for(size_t incoming)
	{
		if (count + incoming > buf.size())
		{
			const size_t drop = count + incoming - buf.size();
			head = (head + drop) % buf.size();
			count -= drop;
		}
	}

	void push(const BYTE* data, size_t bytes, bool silent)
	{
		std::lock_guard<std::mutex> lock(mutex);
		if (bytes > buf.size())
		{
			if (data)
			{
				data += bytes - buf.size();
			}
			bytes = buf.size();
		}
		drop_for(bytes);
		const size_t tail = (head + count) % buf.size();
		const size_t first = std::min(bytes, buf.size() - tail);
		if (silent || !data)
		{
			std::memset(&buf[tail], 0, first);
			if (bytes > first)
			{
				std::memset(&buf[0], 0, bytes - first);
			}
		}
		else
		{
			std::memcpy(&buf[tail], data, first);
			if (bytes > first)
			{
				std::memcpy(&buf[0], data + first, bytes - first);
			}
		}
		count += bytes;
	}

	size_t available() const
	{
		return count;
	}

	// Copy up to `bytes` into `dst`; returns how many bytes were available.
	size_t pop(BYTE* dst, size_t bytes)
	{
		std::lock_guard<std::mutex> lock(mutex);
		bytes = std::min(bytes, count);
		const size_t first = std::min(bytes, buf.size() - head);
		std::memcpy(dst, &buf[head], first);
		if (bytes > first)
		{
			std::memcpy(dst + first, &buf[0], bytes - first);
		}
		head = (head + bytes) % buf.size();
		count -= bytes;
		return bytes;
	}
};

} // namespace

AudioMirror::~AudioMirror()
{
	stop();
}

std::string AudioMirror::status() const
{
	std::lock_guard<std::mutex> lock(status_mutex_);
	return status_;
}

void AudioMirror::set_status(std::string s)
{
	std::lock_guard<std::mutex> lock(status_mutex_);
	status_ = std::move(s);
}

bool AudioMirror::start(DWORD pid)
{
	stop();
	if (!pid)
	{
		set_status("No target process.");
		return false;
	}
	stop_event_ = CreateEventW(nullptr, TRUE, FALSE, nullptr);
	if (!stop_event_)
	{
		set_status("CreateEvent failed.");
		return false;
	}
	pid_ = pid;
	set_status("Starting…");
	thread_ = std::thread(&AudioMirror::thread_main, this, pid);
	return true;
}

void AudioMirror::stop()
{
	if (stop_event_)
	{
		SetEvent(stop_event_);
	}
	if (thread_.joinable())
	{
		thread_.join();
	}
	if (stop_event_)
	{
		CloseHandle(stop_event_);
		stop_event_ = nullptr;
	}
	for (auto& shm : audio_ring_shm_)
	{
		shm.reset();
	}
	running_.store(false, std::memory_order_release);
	source_.store(Source::None, std::memory_order_relaxed);
	buffered_ms_.store(0, std::memory_order_relaxed);
	pid_ = 0;
}

bool AudioMirror::stop_requested() const
{
	return stop_event_ != nullptr && WaitForSingleObject(stop_event_, 0) == WAIT_OBJECT_0;
}

bool AudioMirror::wait_for_format(AudioRingHeader* ring, DWORD timeout_ms)
{
	const DWORD end = GetTickCount() + timeout_ms;
	for (;;)
	{
		if (audio_ring_format_ready(*ring))
		{
			return true;
		}
		if (stop_event_ && WaitForSingleObject(stop_event_, 25) == WAIT_OBJECT_0)
		{
			return false; // stopping
		}
		if (GetTickCount() >= end)
		{
			return false; // hook never published a format -> fall back to loopback
		}
	}
}

void AudioMirror::thread_main(DWORD pid)
{
	const bool com_ok = SUCCEEDED(CoInitializeEx(nullptr, COINIT_MULTITHREADED));

	// Create the shared audio ring the injected hook produces into, and enable
	// capture. If the hook is present it publishes a format within ~1 s and we
	// consume the ring (no echo); otherwise we fall back to process loopback.
	bool handled = false;
	AudioRingHeader* rings[kMaxAudioStreams] = {};
	bool created_primary = false;
	for (unsigned i = 0; i < kMaxAudioStreams; ++i)
	{
		if (audio_ring_shm_[i].create(audio_ring_name(pid, i), audio_ring_total_size(kAudioRingCapacity)))
		{
			rings[i] = audio_ring_shm_[i].as<AudioRingHeader>();
			audio_ring_init(*rings[i], kAudioRingCapacity);
			rings[i]->capture_enabled.store(1, std::memory_order_release);
			created_primary = created_primary || (i == 0);
		}
	}
	if (created_primary)
	{
		set_status("Waiting for render-hook…");
		if (wait_for_format(rings[0], 1000))
		{
			handled = run_hooked(rings);
		}
	}

	if (!handled && !stop_requested())
	{
		run_loopback(pid);
	}

	for (auto& shm : audio_ring_shm_)
	{
		shm.reset();
	}

	if (running_.load(std::memory_order_acquire))
	{
		running_.store(false, std::memory_order_release);
		set_status("Stopped.");
	}
	source_.store(Source::None, std::memory_order_relaxed);

	if (com_ok)
	{
		CoUninitialize();
	}
}

// Consume the render-hook's shared ring and re-render the game's frames. The
// game is silenced locally by the hook, so the operator hears no echo. Returns
// true if it ran to a clean stop; false on setup failure (caller falls back).
bool AudioMirror::run_hooked(AudioRingHeader* const* rings)
{
	AudioRingHeader* primary = rings[0];
	auto disable_all = [&] {
		for (unsigned i = 0; i < kMaxAudioStreams; ++i)
		{
			if (rings[i] != nullptr)
			{
				rings[i]->capture_enabled.store(0, std::memory_order_release);
			}
		}
	};
	auto fail_to_loopback = [&] {
		disable_all(); // let the game play locally again
		return false;
	};

	const unsigned rate = primary->sample_rate;
	const unsigned channels = primary->channels;
	const unsigned bits = primary->bits;
	const unsigned tag = primary->format_tag;
	const unsigned block_align = primary->block_align ? primary->block_align : channels * (bits / 8);
	if (rate == 0 || channels == 0 || block_align == 0)
	{
		return fail_to_loopback();
	}

	// Reconstruct the game's WAVEFORMATEX and let shared-mode WASAPI convert it
	// to the endpoint format via AUTOCONVERTPCM.
	WAVEFORMATEXTENSIBLE wfx = {};
	wfx.Format.nChannels = static_cast<WORD>(channels);
	wfx.Format.nSamplesPerSec = rate;
	wfx.Format.wBitsPerSample = static_cast<WORD>(bits);
	wfx.Format.nBlockAlign = static_cast<WORD>(block_align);
	wfx.Format.nAvgBytesPerSec = block_align * rate;
	if (channels > 2 || bits > 16)
	{
		wfx.Format.wFormatTag = WAVE_FORMAT_EXTENSIBLE;
		wfx.Format.cbSize = sizeof(WAVEFORMATEXTENSIBLE) - sizeof(WAVEFORMATEX);
		wfx.Samples.wValidBitsPerSample = static_cast<WORD>(bits);
		switch (channels) // best-effort default channel masks
		{
		case 6:
			wfx.dwChannelMask = 0x3F;
			break;
		case 8:
			wfx.dwChannelMask = 0xFF;
			break;
		default:
			wfx.dwChannelMask = (channels >= 32) ? 0xFFFFFFFFu : ((1u << channels) - 1u);
			break;
		}
		wfx.SubFormat =
		    (tag == WAVE_FORMAT_IEEE_FLOAT) ? KSDATAFORMAT_SUBTYPE_IEEE_FLOAT : KSDATAFORMAT_SUBTYPE_PCM;
	}
	else
	{
		wfx.Format.wFormatTag = static_cast<WORD>(tag ? tag : WAVE_FORMAT_PCM);
		wfx.Format.cbSize = 0;
	}
	auto* fmt = reinterpret_cast<WAVEFORMATEX*>(&wfx);

	IMMDeviceEnumerator* enumerator = nullptr;
	IMMDevice* endpoint = nullptr;
	IAudioClient* render_client = nullptr;
	IAudioRenderClient* render = nullptr;
	HANDLE render_event = nullptr;
	bool started = false;

	auto fail = [&](const char* msg, HRESULT hr) {
		char buf[160];
		std::snprintf(buf, sizeof(buf), "%s (0x%08lX)", msg, static_cast<unsigned long>(hr));
		set_status(buf);
	};

	do
	{
		HRESULT hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
		                              __uuidof(IMMDeviceEnumerator), reinterpret_cast<void**>(&enumerator));
		if (FAILED(hr))
		{
			fail("CoCreateInstance(MMDeviceEnumerator)", hr);
			break;
		}
		hr = enumerator->GetDefaultAudioEndpoint(eRender, eConsole, &endpoint);
		if (FAILED(hr))
		{
			fail("GetDefaultAudioEndpoint", hr);
			break;
		}
		hr = endpoint->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr,
		                        reinterpret_cast<void**>(&render_client));
		if (FAILED(hr))
		{
			fail("Activate render client", hr);
			break;
		}

		render_event = CreateEventW(nullptr, FALSE, FALSE, nullptr);
		if (!render_event)
		{
			fail("CreateEvent(render)", HRESULT_FROM_WIN32(GetLastError()));
			break;
		}

		constexpr REFERENCE_TIME kRenderBuffer = 30 * 10000; // 30 ms
		const DWORD flags = AUDCLNT_STREAMFLAGS_EVENTCALLBACK | AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM |
		                    AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY;
		hr = render_client->Initialize(AUDCLNT_SHAREMODE_SHARED, flags, kRenderBuffer, 0, fmt, nullptr);
		if (FAILED(hr))
		{
			// The game's format isn't renderable here (rare). Bail to loopback.
			break;
		}
		started = true; // past the point where falling back is clean
		if (FAILED(hr = render_client->SetEventHandle(render_event)))
		{
			fail("Render SetEventHandle", hr);
			break;
		}
		if (FAILED(hr = render_client->GetService(__uuidof(IAudioRenderClient),
		                                          reinterpret_cast<void**>(&render))))
		{
			fail("GetService(RenderClient)", hr);
			break;
		}
		UINT32 render_frames = 0;
		if (FAILED(hr = render_client->GetBufferSize(&render_frames)))
		{
			fail("GetBufferSize", hr);
			break;
		}

		sample_rate_.store(rate, std::memory_order_relaxed);
		channels_.store(channels, std::memory_order_relaxed);

		const size_t frame_bytes = block_align;
		const size_t prime_bytes = frame_bytes * (rate * 30 / 1000); // ~30 ms before feeding
		bool primed = false;

		// Mixing scratch (only used when >1 same-format stream is active): a per-stream
		// temp buffer and a float accumulator sized to the render buffer.
		const bool mixer_ok = mix_format_supported(tag, bits);
		std::vector<BYTE> temp(static_cast<size_t>(render_frames) * frame_bytes);
		std::vector<float> acc(static_cast<size_t>(render_frames) * channels);

		if (FAILED(hr = render_client->Start()))
		{
			fail("Render Start", hr);
			break;
		}

		char st[96];
		std::snprintf(st, sizeof(st), "Mirroring (hooked, no echo): %u Hz %u ch %u-bit", rate, channels, bits);
		set_status(st);
		source_.store(Source::Hooked, std::memory_order_relaxed);
		running_.store(true, std::memory_order_release);

		HANDLE waits[2] = {stop_event_, render_event};
		for (;;)
		{
			const DWORD w = WaitForMultipleObjects(2, waits, FALSE, 200);
			if (w == WAIT_OBJECT_0)
			{
				break; // stop requested
			}

			UINT32 padding = 0;
			if (FAILED(render_client->GetCurrentPadding(&padding)))
			{
				continue;
			}
			const UINT32 avail = render_frames - padding;
			// The primary ring is the master clock for priming + how much to write.
			const std::uint32_t ring_bytes = audio_ring_available(*primary);
			buffered_ms_.store(static_cast<unsigned>(ring_bytes / frame_bytes * 1000 / rate),
			                   std::memory_order_relaxed);
			if (!primed && ring_bytes >= prime_bytes)
			{
				primed = true;
			}
			if (primed && avail > 0)
			{
				const UINT32 have = static_cast<UINT32>(ring_bytes / frame_bytes);
				const UINT32 to_write = std::min(avail, have);
				if (to_write > 0)
				{
					// Active streams = same format as primary (so they can be summed).
					// Streams with a different format are still silenced by the hook (no
					// echo) but can't be mixed here without resampling -> skipped.
					unsigned active[kMaxAudioStreams];
					unsigned n_active = 0;
					for (unsigned i = 0; i < kMaxAudioStreams; ++i)
					{
						AudioRingHeader* r = rings[i];
						if (r == nullptr)
						{
							continue;
						}
						if (i == 0 || (audio_ring_format_ready(*r) && r->sample_rate == rate &&
						               r->channels == channels && r->bits == bits && r->format_tag == tag))
						{
							active[n_active++] = i;
						}
					}

					BYTE* dst = nullptr;
					if (SUCCEEDED(render->GetBuffer(to_write, &dst)))
					{
						const std::uint32_t want_bytes = to_write * static_cast<std::uint32_t>(frame_bytes);
						if (n_active <= 1 || !mixer_ok)
						{
							// Single stream (the common case) or an unmixable format:
							// passthrough the primary, byte-for-byte (no mixer overhead).
							audio_ring_pop(*primary, dst, want_bytes);
						}
						else
						{
							const std::uint32_t samples = to_write * channels;
							std::fill(acc.begin(), acc.begin() + samples, 0.0f);
							for (unsigned k = 0; k < n_active; ++k)
							{
								std::memset(temp.data(), 0, want_bytes); // zero-fill short reads
								audio_ring_pop(*rings[active[k]], temp.data(), want_bytes);
								mix_add(acc.data(), temp.data(), samples, tag, bits);
							}
							mix_store(dst, acc.data(), samples, tag, bits);
						}
						render->ReleaseBuffer(to_write, 0);
					}
				}
				if (to_write < avail)
				{
					primed = false; // ran dry; rebuffer before resuming
				}
			}
		}

		render_client->Stop();
	} while (false);

	disable_all(); // game audible again on stop

	if (render)
	{
		render->Release();
	}
	if (render_client)
	{
		render_client->Release();
	}
	if (endpoint)
	{
		endpoint->Release();
	}
	if (enumerator)
	{
		enumerator->Release();
	}
	if (render_event)
	{
		CloseHandle(render_event);
	}

	if (!started)
	{
		// Never got a working render client; let the caller try loopback. Leave
		// capture disabled (already cleared above) so loopback hears the game.
		return false;
	}
	return true;
}

void AudioMirror::run_loopback(DWORD pid)
{
	source_.store(Source::Loopback, std::memory_order_relaxed);

	IMMDeviceEnumerator* enumerator = nullptr;
	IMMDevice* endpoint = nullptr;
	IAudioClient* render_client = nullptr;
	IAudioRenderClient* render = nullptr;
	WAVEFORMATEX* fmt = nullptr;
	HANDLE render_event = nullptr;
	ProcessLoopbackCapture capture;

	auto fail = [&](const char* msg, HRESULT hr) {
		char buf[160];
		std::snprintf(buf, sizeof(buf), "%s (0x%08lX)", msg, static_cast<unsigned long>(hr));
		set_status(buf);
	};

	do
	{
		HRESULT hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
		                              __uuidof(IMMDeviceEnumerator), reinterpret_cast<void**>(&enumerator));
		if (FAILED(hr))
		{
			fail("CoCreateInstance(MMDeviceEnumerator)", hr);
			break;
		}
		hr = enumerator->GetDefaultAudioEndpoint(eRender, eConsole, &endpoint);
		if (FAILED(hr))
		{
			fail("GetDefaultAudioEndpoint", hr);
			break;
		}
		hr = endpoint->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr,
		                        reinterpret_cast<void**>(&render_client));
		if (FAILED(hr))
		{
			fail("Activate render client", hr);
			break;
		}
		// Capture and render share one format (the output endpoint's mix format);
		// WASAPI converts the captured process audio into it.
		hr = render_client->GetMixFormat(&fmt);
		if (FAILED(hr))
		{
			fail("GetMixFormat", hr);
			break;
		}
		sample_rate_.store(fmt->nSamplesPerSec, std::memory_order_relaxed);
		channels_.store(fmt->nChannels, std::memory_order_relaxed);

		render_event = CreateEventW(nullptr, FALSE, FALSE, nullptr);
		if (!render_event)
		{
			fail("CreateEvent(render)", HRESULT_FROM_WIN32(GetLastError()));
			break;
		}

		constexpr REFERENCE_TIME kRenderBuffer = 30 * 10000; // 30 ms, in 100-ns units
		hr = render_client->Initialize(AUDCLNT_SHAREMODE_SHARED, AUDCLNT_STREAMFLAGS_EVENTCALLBACK,
		                               kRenderBuffer, 0, fmt, nullptr);
		if (FAILED(hr))
		{
			fail("Render Initialize", hr);
			break;
		}
		hr = render_client->SetEventHandle(render_event);
		if (FAILED(hr))
		{
			fail("Render SetEventHandle", hr);
			break;
		}
		hr = render_client->GetService(__uuidof(IAudioRenderClient), reinterpret_cast<void**>(&render));
		if (FAILED(hr))
		{
			fail("GetService(RenderClient)", hr);
			break;
		}
		UINT32 render_frames = 0;
		hr = render_client->GetBufferSize(&render_frames);
		if (FAILED(hr))
		{
			fail("GetBufferSize", hr);
			break;
		}

		const size_t frame_bytes = fmt->nBlockAlign;
		ByteRing ring;
		ring.init(frame_bytes * fmt->nSamplesPerSec); // ~1 s of slack
		// Build ~30 ms of buffer before feeding the renderer, and rebuild it after
		// an underrun, so brief capture gaps don't continuously glitch.
		const size_t prime_bytes = frame_bytes * (fmt->nSamplesPerSec * 30 / 1000);
		bool primed = false;

		// Capture pushes packets straight into the render ring.
		if (!capture.start(pid, fmt, [&ring, frame_bytes](const BYTE* data, UINT32 frames, bool silent) {
			    ring.push(data, static_cast<size_t>(frames) * frame_bytes, silent);
		    }))
		{
			fail("Capture start", E_FAIL);
			break;
		}

		if (FAILED(hr = render_client->Start()))
		{
			fail("Render Start", hr);
			break;
		}

		char st[112];
		std::snprintf(st, sizeof(st), "Mirroring (loopback, echo): device endpoint %u Hz %u ch",
		              static_cast<unsigned>(fmt->nSamplesPerSec), static_cast<unsigned>(fmt->nChannels));
		set_status(st);
		running_.store(true, std::memory_order_release);

		HANDLE waits[2] = {stop_event_, render_event};
		for (;;)
		{
			const DWORD w = WaitForMultipleObjects(2, waits, FALSE, 200);
			if (w == WAIT_OBJECT_0)
			{
				break;
			}
			if (!capture.running())
			{
				set_status(capture.status());
				break;
			}

			UINT32 padding = 0;
			if (FAILED(render_client->GetCurrentPadding(&padding)))
			{
				continue;
			}
			const UINT32 avail = render_frames - padding;
			buffered_ms_.store(
			    static_cast<unsigned>(ring.available() / frame_bytes * 1000 / fmt->nSamplesPerSec),
			    std::memory_order_relaxed);
			if (!primed && ring.available() >= prime_bytes)
			{
				primed = true;
			}
			if (primed && avail > 0)
			{
				const UINT32 have = static_cast<UINT32>(ring.available() / frame_bytes);
				const UINT32 to_write = std::min(avail, have);
				if (to_write > 0)
				{
					BYTE* dst = nullptr;
					if (SUCCEEDED(render->GetBuffer(to_write, &dst)))
					{
						ring.pop(dst, static_cast<size_t>(to_write) * frame_bytes);
						render->ReleaseBuffer(to_write, 0);
					}
				}
				if (to_write < avail)
				{
					primed = false; // ran dry; rebuffer before resuming
				}
			}
		}

		render_client->Stop();
	} while (false);

	capture.stop();

	if (render)
	{
		render->Release();
	}
	if (render_client)
	{
		render_client->Release();
	}
	if (endpoint)
	{
		endpoint->Release();
	}
	if (enumerator)
	{
		enumerator->Release();
	}
	if (fmt)
	{
		CoTaskMemFree(fmt);
	}
	if (render_event)
	{
		CloseHandle(render_event);
	}
}

} // namespace coop