#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <mmsystem.h>
#include <avrt.h>
#include <winternl.h>

#define SPIN_THRESHOLD_NS 1000000	// NOTE(peter): 2000µs spin threshold for Windows

struct timer_handle {
	uint64_t interval_ns;
	uint64_t qpc_frequency;
	uint64_t next_deadline;
	uint32_t running;

	HANDLE event;
	HANDLE timer_thread;
	HANDLE mmcss_handle;

#ifdef TIMER_DEBUG
	uint64_t last_wait_start_ns;
#endif
};

typedef LONG NTSTATUS;
typedef NTSTATUS (__stdcall *NtDelayExecution_t)(BOOLEAN Alertable, PLARGE_INTEGER DelayInterval);

static NtDelayExecution_t pNtDelayExecution;
#define NS_TO_100NS_UNITS 100

static inline uint64_t qpc_now_ns(uint64_t freq) {
	LARGE_INTEGER qpc;
	QueryPerformanceCounter(&qpc);

	uint64_t q = (uint64_t)qpc.QuadPart;
	uint64_t s = q / freq;
	uint64_t r = q % freq;

	return s * 1000000000ULL + (r * 1000000000ULL + (freq >> 1)) / freq;
}

static void try_set_nt_timer_resolution(void) {
	typedef LONG NTSTATUS;
	typedef NTSTATUS (__stdcall *NtSetTimerResolution_t)(ULONG, BOOLEAN, ULONG *);
	NtSetTimerResolution_t NtSetTimerResolution = (NtSetTimerResolution_t)(uintptr_t)GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "NtSetTimerResolution");
	if(NtSetTimerResolution) {
		ULONG requested = 5000; /* 0.5ms = 5000 * 100ns */
		ULONG actual = 0;
		NtSetTimerResolution(requested, TRUE, &actual);
	}
}

static void set_realtime_priority(HANDLE *mmcss_handle) {
	DWORD task_index = 0;
	*mmcss_handle = AvSetMmThreadCharacteristicsW(L"Pro Audio", &task_index);
	if(!*mmcss_handle) {
		SetPriorityClass(GetCurrentProcess(), REALTIME_PRIORITY_CLASS);
		SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);
	}
}

static void timer_sleep(uint64_t ns_to_delay) {
	LONGLONG delay_100ns = -(LONGLONG)(ns_to_delay / NS_TO_100NS_UNITS);
	if(delay_100ns < 0) {
		LARGE_INTEGER delay;
		delay.QuadPart = delay_100ns;
		pNtDelayExecution(FALSE, &delay);
	}
}

static DWORD WINAPI timer_thread_func(LPVOID arg) {
	struct timer_handle *t = (struct timer_handle *)arg;
	SetThreadAffinityMask(GetCurrentThread(), 1);
	set_realtime_priority(&t->mmcss_handle);

	while(t->running) {
#ifdef TIMER_DEBUG
		int64_t remaining_after_sleep_ns = -1;
#endif
		uint64_t now = qpc_now_ns(t->qpc_frequency);

		if(now < t->next_deadline) {
			uint64_t diff = t->next_deadline - now;
			if(diff > SPIN_THRESHOLD_NS) {
				timer_sleep(diff - SPIN_THRESHOLD_NS);
#ifdef TIMER_DEBUG
				now = qpc_now_ns(t->qpc_frequency);
				remaining_after_sleep_ns = (int64_t)(t->next_deadline - now);
#endif
			}
			while(qpc_now_ns(t->qpc_frequency) < t->next_deadline) {
				_mm_pause();
			}
		}

		now = qpc_now_ns(t->qpc_frequency);
		SetEvent(t->event);

#ifdef TIMER_DEBUG
		if(t->last_wait_start_ns > 0) {
			int64_t overshoot_ns = (int64_t)(now - t->next_deadline);
			if(overshoot_ns < 0) overshoot_ns = 0;

			if(remaining_after_sleep_ns >= 0) {
				DEBUG_PRINT("[DEBUG] Woke up with %lld ns left. Overshoot: %5lld ns\n", remaining_after_sleep_ns, overshoot_ns);
			} else {
				DEBUG_PRINT("[DEBUG] No sleep. Overshoot: %lld ns\n", overshoot_ns);
			}
		}
		t->last_wait_start_ns = now;
#endif

		t->next_deadline += t->interval_ns;
	}

	return 0;
}

static void timer_init(void) {
	if (pNtDelayExecution == NULL) {
		pNtDelayExecution = (NtDelayExecution_t)(uintptr_t)GetProcAddress( GetModuleHandleW(L"ntdll.dll"), "NtDelayExecution" );
	}
	timeBeginPeriod(1);
	try_set_nt_timer_resolution();
}

static void timer_shutdown(void) {
	timeEndPeriod(1);
}


static struct timer_handle *timer_new(uint64_t interval_ns) {
	struct timer_handle *t = calloc(1, sizeof(struct timer_handle));

	LARGE_INTEGER freq;
	QueryPerformanceFrequency(&freq);
	t->qpc_frequency = freq.QuadPart;
	t->interval_ns = interval_ns;
	t->next_deadline = qpc_now_ns(t->qpc_frequency) + interval_ns;
	t->running = 1;
	t->mmcss_handle = 0;

#ifdef TIMER_DEBUG
	t->last_wait_start_ns = 0;
#endif

	t->event = CreateEvent(0, FALSE, FALSE, 0);
	t->timer_thread = CreateThread(0, 0, timer_thread_func, t, 0, 0);

	return t;
}

static uint32_t timer_wait(struct timer_handle *t) {
	WaitForSingleObject(t->event, INFINITE);
	return 1;
}

static void timer_destroy(struct timer_handle *t) {
	t->running = 0;

	SetEvent(t->event);
	WaitForSingleObject(t->timer_thread, INFINITE);
	CloseHandle(t->timer_thread);
	CloseHandle(t->event);

	if(t->mmcss_handle) {
		AvRevertMmThreadCharacteristics(t->mmcss_handle);
	}

	free(t);
}