#include "muscles.h"
#include "structs.h"
#include "search.h"

#include <cstdint>
#include <algorithm>

static void *thread = nullptr;
static bool started = false;
static bool running = false;

static u64 *results = nullptr;
static int n_results = 0;

static u64 *prev_results = nullptr;
static int n_prev_results = 0;

static Search search;
static std::vector<std::pair<u64, u64>> ranges;

void perform_search();

void start_search(Search& s, std::vector<Region> const& regions) {
	if (running)
		return;

	ranges.resize(0);
	for (auto& reg : regions)
		ranges.push_back(std::make_pair<u64, u64>((u64)reg.base, (u64)reg.size));

	std::sort(ranges.begin(), ranges.end(), [](std::pair<u64, u64>& a, std::pair<u64, u64>& b) {
		return a.first < b.first;
	});

	search.params = s.params;
	if (search.params) {
		search.n_params = s.n_params;
		search.record = s.record;
	}
	else
		search.single_value = s.single_value;

	search.byte_align = s.byte_align;

	search.start_addr = s.start_addr;
	search.end_addr = s.end_addr;

	search.source_type = s.source_type;
	search.pid = s.pid;
	search.identifier = s.identifier;

	auto func = [](void *data) {
		perform_search();
		return (THREAD_RETURN_TYPE)0;
	};
	start_thread(&thread, nullptr, func);
}

bool check_search_running() {
	return running;
}

bool check_search_finished() {
	bool finished = started && !running;
	if (finished)
		started = false;

	return finished;
}

void get_search_results(std::vector<u64>& results_vec) {
	if (!results)
		return;

	results_vec.resize(n_results);
	if (n_results)
		memcpy(results_vec.data(), results, n_results * sizeof(u64));
}

void reset_search() {
	n_results = 0;
}

void exit_search() {
	if (results) {
		delete[] results;
		results = nullptr;
	}
	if (prev_results) {
		delete[] prev_results;
		prev_results = nullptr;
	}
}

struct Scan_Range {
	u64 start;
	int first_range;
	int last_range;
};

Scan_Range isolate_scan_ranges() {
	Scan_Range scan = {
		.start = search.start_addr,
		.first_range = -1,
		.last_range = -1
	};

	for (auto& r : ranges) {
		scan.first_range++;
		if (scan.start < r.first + r.second) {
			if (scan.start < r.first)
				scan.start = r.first;
			break;
		}
	}
	for (auto& r : ranges) {
		if (search.end_addr < r.first)
			break;

		scan.last_range++;
	}

	return scan;
}

template <int method, typename T>
void single_value_search(SOURCE_HANDLE handle, T v1, T v2) {
	int byte_align = search.byte_align;
	if (byte_align <= 0)
		byte_align = sizeof(T);

	char *buf = new char[PAGE_SIZE]();

	if (n_results == 0) {
		auto scan = isolate_scan_ranges();

		u64 addr = scan.start;
		for (int i = scan.first_range; i <= scan.last_range && addr <= search.end_addr; i++) {
			u64 range_end = ranges[i].first + ranges[i].second;
			if (search.end_addr < range_end)
				range_end = search.end_addr;

			u64 page = addr & ~(PAGE_SIZE - 1);
			int offset = (int)(addr - page) & ~(sizeof(T) - 1);

			for (; page < range_end; page += PAGE_SIZE) {
				int retrieved = read_page(handle, search.source_type, page, buf);
				if (retrieved <= 0)
					continue;

				for (int j = offset; j <= PAGE_SIZE - sizeof(T); j += byte_align) {
					T value = *(T*)(&buf[j]);
					if constexpr (method == METHOD_EQUALS) {
						if (value == v1) {
							results[n_results++] = page + (u64)j;
							if (n_results >= MAX_SEARCH_RESULTS)
								goto done_single;
						}
					}
					else {
						if (value >= v1 && value <= v2) {
							results[n_results++] = page + (u64)j;
							if (n_results >= MAX_SEARCH_RESULTS)
								goto done_single;
						}
					}
				}
			}

			if (i < scan.last_range)
				addr = ranges[i + 1].first;
		}
	}
	else {
		n_prev_results = n_results;
		if (!prev_results)
			prev_results = new u64[MAX_SEARCH_RESULTS];

		memcpy(prev_results, results, n_prev_results * sizeof(u64));
		n_results = 0;

		u64 page = 0;
		bool fail = false;

		for (int i = 0; i < n_prev_results && n_results < MAX_SEARCH_RESULTS; i++) {
			u64 addr = prev_results[i];
			u64 p = addr & ~(PAGE_SIZE - 1);
			int offset = (int)(addr - p);//(int)(addr & (u64)(PAGE_SIZE - 1));

			if (i == 0 || p > page) {
				page = p;
				int retrieved = read_page(handle, search.source_type, page, buf);
				fail = retrieved <= 0;
			}
			if (!fail) {
				T value = *(T*)(&buf[offset]);
				if constexpr (method == METHOD_EQUALS) {
					if (value == v1)
						results[n_results++] = addr;
				}
				else {
					if (value >= v1 && value <= v2)
						results[n_results++] = addr;
				}
			}
		}
	}

done_single:
	delete[] buf;
}

// DRY: Do Repeat Yourself
void do_single_value_search(SOURCE_HANDLE handle) {
	Search_Parameter sv = search.single_value;
	u32 flags = sv.flags & FIELD_FLAGS;

	if (sv.method == METHOD_EQUALS) {
		if (flags & FLAG_FLOAT) {
			if (sv.size == 32)
				single_value_search<METHOD_EQUALS, float>(handle, *(float*)&sv.value1, 0);
			else if (sv.size == 64)
				single_value_search<METHOD_EQUALS, double>(handle, *(double*)&sv.value1, 0);
		}
		else if (flags & FLAG_SIGNED) {
			if (sv.size == 8)
				single_value_search<METHOD_EQUALS, std::int8_t>(handle, *(std::int8_t*)&sv.value1, 0);
			else if (sv.size == 16)
				single_value_search<METHOD_EQUALS, std::int16_t>(handle, *(std::int16_t*)&sv.value1, 0);
			else if (sv.size == 32)
				single_value_search<METHOD_EQUALS, std::int32_t>(handle, *(std::int32_t*)&sv.value1, 0);
			else
				single_value_search<METHOD_EQUALS, std::int64_t>(handle, *(std::int64_t*)&sv.value1, 0);
		}
		else {
			if (sv.size == 8)
				single_value_search<METHOD_EQUALS, std::uint8_t>(handle, *(std::uint8_t*)&sv.value1, 0);
			else if (sv.size == 16)
				single_value_search<METHOD_EQUALS, std::uint16_t>(handle, *(std::uint16_t*)&sv.value1, 0);
			else if (sv.size == 32)
				single_value_search<METHOD_EQUALS, std::uint32_t>(handle, *(std::uint32_t*)&sv.value1, 0);
			else
				single_value_search<METHOD_EQUALS, std::uint64_t>(handle, *(std::uint64_t*)&sv.value1, 0);
		}
	}
	else if (sv.method == METHOD_RANGE) {
		if (flags & FLAG_FLOAT) {
			if (sv.size == 32)
				single_value_search<METHOD_RANGE, float>(handle, *(float*)&sv.value1, *(float*)&sv.value2);
			else if (sv.size == 64)
				single_value_search<METHOD_RANGE, double>(handle, *(double*)&sv.value1, *(double*)&sv.value2);
		}
		else if (flags & FLAG_SIGNED) {
			if (sv.size == 8)
				single_value_search<METHOD_RANGE, std::int8_t>(handle, *(std::int8_t*)&sv.value1, *(std::int8_t*)&sv.value2);
			else if (sv.size == 16)
				single_value_search<METHOD_RANGE, std::int16_t>(handle, *(std::int16_t*)&sv.value1, *(std::int16_t*)&sv.value2);
			else if (sv.size == 32)
				single_value_search<METHOD_RANGE, std::int32_t>(handle, *(std::int32_t*)&sv.value1, *(std::int32_t*)&sv.value2);
			else
				single_value_search<METHOD_RANGE, std::int64_t>(handle, *(std::int64_t*)&sv.value1, *(std::int64_t*)&sv.value2);
		}
		else {
			if (sv.size == 8)
				single_value_search<METHOD_RANGE, std::uint8_t>(handle, *(std::uint8_t*)&sv.value1, *(std::uint8_t*)&sv.value2);
			else if (sv.size == 16)
				single_value_search<METHOD_RANGE, std::uint16_t>(handle, *(std::uint16_t*)&sv.value1, *(std::uint16_t*)&sv.value2);
			else if (sv.size == 32)
				single_value_search<METHOD_RANGE, std::uint32_t>(handle, *(std::uint32_t*)&sv.value1, *(std::uint32_t*)&sv.value2);
			else
				single_value_search<METHOD_RANGE, std::uint64_t>(handle, *(std::uint64_t*)&sv.value1, *(std::uint64_t*)&sv.value2);
		}
	}
}

template<typename T>
s64 roundtrip_cast(s64 in) {
	T value = (T)in;
	return (s64)value;
}

s64 cast(u32 flags, int size, s64 in) {
	s64 out = 0;

	if (flags & FLAG_FLOAT) {
		out = in;
	}
	else if (flags & FLAG_SIGNED) {
		if (size == 8)
			out = roundtrip_cast<std::int8_t>(in);
		else if (size == 16)
			out = roundtrip_cast<std::int16_t>(in);
		else if (size == 32)
			out = roundtrip_cast<std::int32_t>(in);
		else if (size == 64)
			out = roundtrip_cast<std::int64_t>(in);
		else
			out = in;
	}
	else {
		if (size == 8)
			out = roundtrip_cast<std::uint8_t>(in);
		else if (size == 16)
			out = roundtrip_cast<std::uint16_t>(in);
		else if (size == 32)
			out = roundtrip_cast<std::uint32_t>(in);
		else if (size == 64)
			out = roundtrip_cast<std::uint64_t>(in);
		else
			out = in;
	}

	return out;
}

// TODO: Support both endians, bitfields, arrays (including string literals)
void do_object_search(SOURCE_HANDLE handle) {
	int byte_inc = search.byte_align;
	if (byte_inc <= 0)
		byte_inc = search.record->total_size / 8;

	int n_params = search.n_params;

	char *page_attrs_buf = new char[n_params * (2 * sizeof(s64) + sizeof(u64) + sizeof(int))]();

	s64 *values      = reinterpret_cast<s64*>(page_attrs_buf);
	u64 *page_addrs  = reinterpret_cast<u64*>(page_attrs_buf + n_params * 2 * sizeof(s64));
	int *page_idxs   = reinterpret_cast<int*>(page_attrs_buf + n_params * (2 * sizeof(s64) + sizeof(u64)));

	char *page_mem = new char[(n_params + 1) * PAGE_SIZE];

	// totally unnecessary but totally swag
	char *pages = reinterpret_cast<char*>((reinterpret_cast<u64>(page_mem) + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1));

	for (int i = 0; i < n_params; i++) {
		page_addrs[i] = -1;

		Search_Parameter& p = search.params[i];
		values[i*2]   = cast(p.flags, p.size, p.value1);
		values[i*2+1] = cast(p.flags, p.size, p.value2);
	}

	auto search_all_parameters = [&](u64 head_address) {
		int matches = 0;

		for (int j = 0; j < n_params; j++) {
			u64 byte_offset = (u64)(search.params[j].offset / 8);
			u64 addr = head_address + byte_offset;
			u64 page = addr & ~(PAGE_SIZE - 1);
			int page_offset = (int)(addr - page);

			if (page != page_addrs[j]) {
				int p = 0;
				for (; p < n_params; p++) {
					if (page_addrs[p] == page)
						break;
				}

				if (p >= n_params) {
					int retrieved = read_page(handle, search.source_type, page, &pages[j * PAGE_SIZE]);
					if (retrieved <= 0)
						return -1;

					p = j;
				}

				page_addrs[j] = page;
				page_idxs[j] = p;
			}

			char *buf = &pages[page_idxs[j] * PAGE_SIZE];
			int byte_size = search.params[j].size / 8;

			s64 value = 0;
			for (int k = 0; k < byte_size; k++)
				value |= (s64)(buf[page_offset + k] & 0xff) << (s64)(8*k);

			value = cast(search.params[j].flags, search.params[j].size, value);

			if (
				(search.params[j].method == METHOD_EQUALS && value == values[2*j]) || 
				(search.params[j].method == METHOD_RANGE  && value >= values[2*j] && value <= values[2*j+1])
			)
				matches++;
		}

		return matches;
	};

	if (n_results == 0) {
		auto scan = isolate_scan_ranges();

		u64 head = scan.start;
		for (int i = scan.first_range; i <= scan.last_range && head <= search.end_addr; i++) {
			u64 range_end = ranges[i].first + ranges[i].second;
			if (search.end_addr < range_end)
				range_end = search.end_addr;

			for (; head < range_end; head += byte_inc) {
				int matches = search_all_parameters(head);

				if (matches == n_params) {
					results[n_results++] = head;
					if (n_results >= MAX_SEARCH_RESULTS)
						goto done_object;
				}
			}

			if (i < scan.last_range)
				head = ranges[i + 1].first;
		}
	}
	else {
		n_prev_results = n_results;
		if (!prev_results)
			prev_results = new u64[MAX_SEARCH_RESULTS];

		memcpy(prev_results, results, n_prev_results * sizeof(u64));
		n_results = 0;

		for (int i = 0; i < n_prev_results; i++) {
			u64 head = prev_results[i];
			int matches = search_all_parameters(head);

			if (matches == n_params) {
				results[n_results++] = head;
				if (n_results >= MAX_SEARCH_RESULTS)
					goto done_object;
			}
		}
	}

done_object:
	delete[] page_attrs_buf;
	delete[] page_mem;
}

// We know the thread has ended if started == true and running == false
void perform_search() {
	started = true;
	running = true;

	if (!results)
		results = new u64[MAX_SEARCH_RESULTS];

	auto handle = (SOURCE_HANDLE)0;
	if (search.source_type == SourceFile)
		handle = get_readonly_file_handle(search.identifier);
	else if (search.source_type == SourceProcess)
		handle = get_readonly_process_handle(search.pid);

	if (!handle) {
		running = false;
		return;
	}

	if (!search.params)
		do_single_value_search(handle);
	else
		do_object_search(handle);

	running = false;
}