#include "asset-manager.hpp"
#include "audio-impl-shared.hpp"
#include "audio.hpp"
#include "ecs.hpp"
#include "game-settings.hpp"
#include "math-helpers.hpp"
#include "pk.h"
#include "window.hpp"
#include <chrono>

#ifdef PKE_AUDIO_IMPL_PIPEWIRE
#include "audio-impl-pw.hpp"
#endif

struct pke_audio_master pke_audio_mstr{};

void pke_audio_process_frames(int64_t frame_count);

void pke_audio_init() {
	pke_audio_mstr.master_volume = 1.f;
	pke_audio_mstr.channel_count = 2;
	// pke_audio_mstr.last_tick_tp = {};
	pke_audio_mstr.last_tick_tp = pkeSettings.steadyClock.now();
	pke_audio_mstr.bkt= pk_mem_bucket_create("pke_audio", PK_MEM_DEFAULT_BUCKET_SIZE, PK_MEMBUCKET_FLAG_NONE);
	for (uint8_t i = 0; i < pke_audio_mstr.channel_count; ++i) {
		pke_audio_mstr.source_volumes[i] = 1.f;
	}
	pke_audio_mstr.buffer_size = PKE_AUDIO_BUFFER_FRAMES * pke_audio_mstr.channel_count * sizeof(float);
	pke_audio_mstr.buffer_frames = 0;
	pke_audio_mstr.elapsed_ns = 0;
	pke_audio_mstr.buffer = pk_new_arr<float>(pke_audio_mstr.buffer_size, pke_audio_mstr.bkt);
#ifdef PKE_AUDIO_IMPL_PIPEWIRE
	pke_audio_pw_init();
#endif
}

void pke_audio_teardown() {
#ifdef PKE_AUDIO_IMPL_PIPEWIRE
	pke_audio_pw_teardown();
#endif
	pke_audio_stop_all();
	pk_arr_reset(&pke_audio_mstr.playing_objects);
	pk_delete_arr<float>(pke_audio_mstr.buffer, pke_audio_mstr.buffer_size, pke_audio_mstr.bkt);
	pk_mem_bucket_destroy(pke_audio_mstr.bkt);

	pke_audio_mstr.bkt = nullptr;
	pke_audio_mstr.master_volume = 0;
	memset(&pke_audio_mstr.source_volumes, 0, sizeof(pke_audio_mstr.source_volumes[0]) * pke_audio_source_T_MAX);
	pke_audio_mstr.channel_count = 0;
	pke_audio_mstr.buffer = nullptr;
	pke_audio_mstr.buffer_frames = 0;
	pke_audio_mstr.elapsed_ns = 0;
	pke_audio_mstr.last_tick_tp = {};

}

void pke_audio_tick(double delta) {
	(void)delta;
	uint32_t i, k;
#ifdef PKE_AUDIO_IMPL_PIPEWIRE
	if (pke_audio_pw.is_needing_output_remapped == true) {
		pke_audio_pw_remap_outputs();
		pke_audio_mstr.last_tick_tp = pkeSettings.steadyClock.now();
		return;
	}
#endif
	pke_audio_mstr.mtx_buffer.lock();
	for (i = 0; i < pke_audio_mstr.playing_objects.next; ++i) {
		for (k = 0; k < pke_audio_mstr.playing_objects[i].play_count; ++k) {
			if (pke_audio_mstr.playing_objects[i].instance_handle[k] == InstanceHandle_MAX) {
				continue;
			}
			CompInstance *inst = ECS_GetInstance(pke_audio_mstr.playing_objects[i].instance_handle[k]);
			if (inst == nullptr || inst->instanceHandle == InstanceHandle_MAX) {
				continue;
			}
			BulletToGlm(
				inst->bt.rigidBody->getWorldTransform().getOrigin(),
				pke_audio_mstr.playing_objects[i].position_source[k]
			);
		}
	}

	constexpr int64_t ns_per_bit_chunk = std::chrono::nanoseconds::period::den / 4000;
	constexpr int64_t bits_per_chunk = PKE_AUDIO_BITRATE / 4000;

	std::chrono::time_point<std::chrono::steady_clock> tp_now = pkeSettings.steadyClock.now();
	auto diff = (tp_now - pke_audio_mstr.last_tick_tp).count();
	pke_audio_mstr.elapsed_ns += diff;

	int64_t frame_count = floor((pke_audio_mstr.elapsed_ns / ns_per_bit_chunk)) * bits_per_chunk;
	if (frame_count > 0) {
		// fprintf(stdout, "[pw_tick] frame_count: %li, running_count: %li, elapsed_ns: %li\n", frame_count, pke_audio_mstr.buffer_frames, pke_audio_mstr.elapsed_ns);
		pke_audio_process_frames(frame_count);
		pke_audio_mstr.elapsed_ns %= ns_per_bit_chunk;
	}
	pke_audio_mstr.last_tick_tp = tp_now;
	pke_audio_mstr.mtx_buffer.unlock();
}

float pke_audio_get_volume(pke_audio_source source) {
	return pke_audio_mstr.source_volumes[static_cast<pke_audio_source_T>(source)];
}

void  pke_audio_set_volume(pke_audio_source source, float volume) {
	pke_audio_mstr.source_volumes[static_cast<pke_audio_source_T>(source)] = volume;
}

bool pke_audio_playing_objects_find_first_by_key(void *user_data, void *arr_data) {
	std::tuple<AssetHandle, pke_audio_source> &tup = *reinterpret_cast<std::tuple<AssetHandle, pke_audio_source>*>(user_data);
	pke_audio_obj &audio_obj = *reinterpret_cast<pke_audio_obj*>(arr_data);
	return std::get<0>(tup) == audio_obj.handle && std::get<1>(tup) == audio_obj.source;
}

void pke_audio_play(AssetHandle handle, pke_audio_source audio_source, pke_audio_flags flags, glm::vec3 position_source, InstanceHandle instance_handle) {
	// TODO validation: audio length (does it fit in uint32_t), etc
	// TODO rethink threading: first-pass only mutex
	std::tuple<AssetHandle, pke_audio_source> tup {handle, audio_source};
	pke_audio_mstr.mtx_buffer.lock();
	int64_t idx = (int64_t)pk_arr_find_first_index(&pke_audio_mstr.playing_objects, &tup, pke_audio_playing_objects_find_first_by_key);
	pke_audio_obj *aobj = NULL;
	if (idx == uint32_t(-1)) {
		AM_Get(handle); // keep the asset in memory, freed when play_count hits 0
		pk_arr_append_t(&pke_audio_mstr.playing_objects, {});
		aobj = &pke_audio_mstr.playing_objects[pke_audio_mstr.playing_objects.next-1];
		memset(aobj, 0, sizeof(pke_audio_obj));
		aobj->handle = handle;
		aobj->source = audio_source;
		aobj->play_count = 1;
		idx = 0;
	} else {
		aobj = &pke_audio_mstr.playing_objects[idx];
		idx = aobj->play_count;
		aobj->play_count += 1;
	}
	assert(idx >= 0 && idx < PKE_AUDIO_MAX_CONCURRENT_COUNT);
	aobj->instance_handle[idx] = instance_handle;
	aobj->position_source[idx] = position_source;
	aobj->flags[idx]           = flags;
	aobj->play_heads[idx]      = 0;
	pke_audio_mstr.mtx_buffer.unlock();
}

void pke_audio_stop_all() {
	// TODO fade-out instead of hard-cut? Maybe that should be a separate function.
	uint32_t u;
	pke_audio_mstr.mtx_buffer.lock();
	for (u = 0; u < pke_audio_mstr.playing_objects.next; ++u) {
		AM_Release(pke_audio_mstr.playing_objects[u].handle);
	}
	pk_arr_clear(&pke_audio_mstr.playing_objects);
	pke_audio_mstr.mtx_buffer.unlock();
}

void pke_audio_process_frames(int64_t frame_count) {
	uint8_t pc, pc2;
	uint32_t i, ii, c;
	int64_t i_frame;
	float *dst;
	float val, vol, vol2, dot, dot2;
	glm::vec3 listener_origin;
	glm::vec3 audio_dir = glm::vec3(0);
	float                               *spatial_volumes = nullptr;
	glm::vec3                           *spatial_normals = nullptr;
	pke_audio_fx_params_reverb          *params_reverb = nullptr;
	pke_audio_fx_params_delay           *params_delay = nullptr;
	pke_audio_fx_params_low_pass_filter *params_low_pass_filter = nullptr;

	int64_t stride = pke_audio_mstr.channel_count * sizeof(float);
	frame_count = PK_MIN(frame_count, (pke_audio_mstr.buffer_size / stride) - pke_audio_mstr.buffer_frames);

	// init
	spatial_volumes        = pk_new_arr<float>(pke_audio_mstr.channel_count, pkeSettings.mem_bkt.game_transient);
	spatial_normals        = pk_new_arr<glm::vec3>(pke_audio_mstr.channel_count, pkeSettings.mem_bkt.game_transient);
	params_reverb          = pk_new_arr<pke_audio_fx_params_reverb>(pke_audio_mstr.channel_count, pkeSettings.mem_bkt.game_transient);
	params_delay           = pk_new_arr<pke_audio_fx_params_delay>(pke_audio_mstr.channel_count, pkeSettings.mem_bkt.game_transient);
	params_low_pass_filter = pk_new_arr<pke_audio_fx_params_low_pass_filter>(pke_audio_mstr.channel_count, pkeSettings.mem_bkt.game_transient);

	dst = pke_audio_mstr.buffer + (pke_audio_mstr.buffer_frames * pke_audio_mstr.channel_count);

	// calculate spatial_normals
	// TODO maybe don't use UBO
	// TODO project-defined?
	// TODO instanced audio manager? (easier to handle explicit stereo sources?)
	listener_origin = glm::inverse(UBO.model) * glm::vec4(0, 0, 0, 1);
	for (c = 0; c < pke_audio_mstr.channel_count; ++c) {
		switch (c) {
			case 0:
				// left-speaker
				if (pke_audio_mstr.channel_count > 2) {
					spatial_normals[c] = glm::normalize(glm::vec3(-1.f,  0.f,  1.f));
				} else {
					spatial_normals[c] = glm::normalize(glm::vec3(-1.f,  0.f,  0.f));
				}
				break;
			case 1:
				// right
				if (pke_audio_mstr.channel_count > 2) {
					spatial_normals[c] = glm::normalize(glm::vec3( 1.f,  0.f,  1.f));
				} else {
					spatial_normals[c] = glm::normalize(glm::vec3( 1.f,  0.f,  0.f));
				}
				break;
			case 2:
				// center
				spatial_normals[c] = glm::normalize(glm::vec3( 0.f,  0.f,  1.f));
				break;
			case 3:
				// subwoofer
				spatial_normals[c] = glm::normalize(glm::vec3( 0.f,  0.f,  0.f));
				break;
			case 4:
				// left surround
				spatial_normals[c] = glm::normalize(glm::vec3(-1.f,  0.f,  0.f));
				break;
			case 5:
				// right surround
				spatial_normals[c] = glm::normalize(glm::vec3( 1.f,  0.f,  0.f));
				break;
			case 6:
				// rear left surround
				spatial_normals[c] = glm::normalize(glm::vec3(-1.f,  0.f, -1.f));
				break;
			case 7:
				// rear right surround
				spatial_normals[c] = glm::normalize(glm::vec3( 1.f,  0.f, -1.f));
				break;
			default:
				break;
		}
	}

	if (pke_audio_mstr.playing_objects.next == 0) {
		memset(dst, 0, sizeof(float) * frame_count * pke_audio_mstr.channel_count);
	}

	// calculate
	for (i = 0; i < pke_audio_mstr.playing_objects.next; ++i) {
		pke_audio_obj *aobj = &pke_audio_mstr.playing_objects[i];
		const Asset *a = AM_Get(aobj->handle);
		vol = 1.0;
		vol *= pke_audio_mstr.master_volume;
		vol *= pke_audio_mstr.source_volumes[static_cast<pke_audio_source_T>(aobj->source)];

		for (pc = 0; pc < aobj->play_count; ++pc) {
			// TODO configurable
			// >= 5.0 meters is 100% volume
			// <=50.0 meters is   0% volume
			float distance_volume = glm::distance(aobj->position_source[pc], listener_origin);
			distance_volume = 1.0 - ((distance_volume - 5.0) / 50 - 5);
			distance_volume = glm::clamp(distance_volume, 0.f, 1.f);

			// calculate panning
			// TODO handle subwoofer explicitly
			// float sum = 0.0;
			for (c = 0; c < pke_audio_mstr.channel_count; ++c) {
				audio_dir = glm::normalize(aobj->position_source[pc] - listener_origin);
				for (ii = 0; ii < 3; ++ii) {
					if (glm::isnan(audio_dir[ii])) {
						audio_dir = glm::vec3(0);
						spatial_volumes[c] = 1.f;
						break;
					}
				}
				dot  = glm::dot(glm::vec3(0.f, 0.f, 1.f), audio_dir);
				dot2 = glm::dot(spatial_normals[c], audio_dir);

				params_low_pass_filter[c].cutoff_freq = 0.f;
				params_reverb[c].reverb_strength = 0.f; // TODO
				params_delay[c].delay_frames = 0; // TODO
				spatial_volumes[c] = PK_CLAMP(dot2, 0.f, 1.f);
				if (pke_audio_mstr.channel_count == 2) {
					/* 2025-07-15 JCB
					 * I changed this to always calculate for "forward".
					 * I believe this is only needed if channel_count == 2; this is to simulate surround through stereo (headphones).
					 * If the user is actually using surround sound, their own ears will do this.
					 */
					if (!glm::isnan(dot)) {
						val  = (dot + 1.f) / 2.f; // padding lerp val
						spatial_volumes[c] = PK_CLAMP(dot2 + lerp(0.5, 1.f, val) + 0.05, 0.0f, 1.f);
						if (dot < 0.0) {
							// 20k max, 500 min
							params_low_pass_filter[c].cutoff_freq = log_interp(40000.f, 500.f, abs(dot));
						}
					}
				}
				/*
				fprintf(stdout, "[pw] normal: % 02.0f,% 02.0f, audio_dir: % 02.6f,% 02.6f, dot: % 02.6f, dot2: % 02.6f, dot3: % 02.6f, cutoff: % 5.6f, spatial_volumes:% 02.6f\n",
					spatial_normals[c].x, spatial_normals[c].z,
					audio_dir.x, audio_dir.z,
					dot, dot2, dot + dot2, params_low_pass_filter[c].cutoff_freq, spatial_volumes[c]);
				*/

				if (isnan(spatial_volumes[c])) {
					/*
					 * || spatial_volumes[c] == 0.0f
					fprintf(stderr, "[pw] NaN or 0: chan: %i, norm: %f,%f,%f, src: %f,%f,%f, origin: %f,%f,%f\n",
						c,
						audio_dir.x, audio_dir.y, audio_dir.z,
						aobj->position_source[pc].x, aobj->position_source[pc].y, aobj->position_source[pc].z,
						listener_origin.x, listener_origin.y, listener_origin.z
						);
					*/
					spatial_volumes[c] = 1.0;
				}

				// sum += spatial_volumes[c];
			}
			// normalize (TODO do we want this?)
			/*
			if (sum > 1.f) {
				for (c = 0; c < pke_audio_mstr.channel_count; ++c) {
					spatial_volumes[c] /= sum;
				}
			}
			*/

			/*
			for (c = 0; c < pke_audio_mstr.channel_count; ++c) {
				fprintf(stderr, "[pw] obj: %i, chan: %i vol: %f, spatial: %f, dist: %f\n",
					i, c, vol, spatial_volumes[c], distance_volume
				);
			}
			*/

			for (i_frame = 0; i_frame < frame_count; ++i_frame) {
				dst = pke_audio_mstr.buffer + (pke_audio_mstr.buffer_frames * pke_audio_mstr.channel_count);
				int64_t buffer_idx_advance = (i_frame * (int64_t)pke_audio_mstr.channel_count);
				// fprintf(stdout, "[pw] frame_count: %li, buffer_idx_advance : %li\n", frame_count, buffer_idx_advance);
				dst += buffer_idx_advance; // advance

				for (c = 0; c < pke_audio_mstr.channel_count; ++c) {
					if (i == 0 && pc == 0) { *dst = 0.f; } // clear buffer as we go

					vol2 = vol;
					if (PK_HAS_FLAG(aobj->flags[pc], pke_audio_flag_pos_spatial)) {
						vol2 *= (spatial_volumes[c] * distance_volume);
					}
					if (vol2 <= 0.0) {
						// fprintf(stderr, "[pw] chan: %i vol2 is <= 0.0\n", c);
						dst += 1;
						continue;
					}

					/*
					if (isnan(vol2)) {
						fprintf(stderr, "[pw] vol2 is NaN\n");
					}
					if (vol2 == 0.0) {
						fprintf(stderr, "[pw] vol2 is 0, %f, %f, %f\n", pke_audio_mstr.master_volume, pke_audio_mstr.source_volumes[c], spatial_volumes[c]);
					}
					*/

					// val = ((float*)a->ptr)[aobj->play_heads[pc]]; // val is read inside fx_low_pass_filter
					val  = pke_audio_fx_low_pass_filter((float*)a->ptr, a->size / sizeof(float), (uint32_t)aobj->play_heads[pc], &params_low_pass_filter[c]);
					val += pke_audio_fx_reverb((float*)a->ptr, a->size / sizeof(float), (uint32_t)aobj->play_heads[pc], &params_reverb[c]);
					val += pke_audio_fx_delay((float*)a->ptr, a->size / sizeof(float), (uint32_t)aobj->play_heads[pc], &params_delay[c]);
					*dst += val * vol2;
					/*
					if (isnan(*dst)) {
						fprintf(stderr, "[pw] *dst is NaN\n");
					}
					*/
					dst += 1;

				}

				// TODO type-specific attributes for assets so we can pre-calculate this or just KNOW the frame length ahead of time
				aobj->play_heads[pc] += 1;
				if (aobj->play_heads[pc] >= a->size / sizeof(float)) {
					for (pc2 = 0; pc + pc2 + 1 < aobj->play_count; ++pc2) {
						aobj->instance_handle[pc+pc2] = aobj->instance_handle[pc+pc2+1];
						aobj->position_source[pc+pc2] = aobj->position_source[pc+pc2+1];
						aobj->flags[pc+pc2]           = aobj->flags[pc+pc2+1];
						aobj->play_heads[pc+pc2]      = aobj->play_heads[pc+pc2+1];
					}
					pc -= 1;
					aobj->play_count -= 1;
					break;
				}
			}
		}

		AM_Release(aobj->handle);
		if (aobj->play_count == 0) {
			AM_Release(aobj->handle);
			pk_arr_remove_at(&pke_audio_mstr.playing_objects, i);
			i -= 1;
		}
	}

	pke_audio_mstr.buffer_frames += frame_count;
}