#include "shell.h"
#include "parser.h"
#include "commands.h"
#include "errors.h"

#include <fmt/core.h>
#include <endian.h>

//---
// afh
//---

static bool matches(char const *data, char const *reference,
	char const *pattern, size_t size)
{
	for(size_t i = 0; i < size; i++) {
		if(pattern[i] && data[i] != reference[i]) return false;
	}
	return true;
}

static int hexa(int c) {
	if(c >= '0' && c <= '9')
		return c - '0';
	return (c | 0x20) - 'a' + 10;
}

struct _afh_args {
	/* String to find */
	std::string reference;
	/* Bytes to match within reference (not all are required) */
	std::string pattern;
	/* Size of search */
	size_t size;
	/* Regions to search in, may be empty */
	std::vector<MemoryRegion> regions;
	/* Distance to show hexump around a match */
	int distance;
	/* Required alignment for matches to be used */
	int align;
};

static _afh_args parse_afh(Session &session, Parser &parser)
{
	_afh_args args;
	args.distance = -1;
	args.align = -1;

	parser.option('a', [&args](std::string const &value){
		args.align = atoi(value.c_str());
	});
	parser.option('d', [&args](std::string const &value){
		args.distance = atoi(value.c_str());
	});

	parser.accept_options();
	std::string needle = parser.str();

	/* Check the structure of the needle */

	if(needle.size() == 0 || needle.size() % 2 != 0)
		throw CommandError("search pattern '{}' should be of even non-zero "
			"size", needle);

	size_t bad_index = needle.find_first_not_of("0123456789abcdefABCDEF.");
	if(bad_index != std::string::npos)
		throw CommandError("invalid character '{}' in seach pattern",
			needle[bad_index]);

	for(size_t i = 0; i < needle.size(); i += 2) {
		char c1 = needle[i], c2 = needle[i+1];
		if((c1 == '.') != (c2 == '.'))
			throw CommandError("invalid search byte '{}{}', should be either "
				"'..' or fully specified", c1, c2);
	}

	/* Convert into a reference/pattern form */

	args.size = needle.size() / 2;
	args.reference.reserve(args.size);
	args.pattern.reserve(args.size);

	for(size_t i = 0; i < args.size; i++) {
		char c1 = needle[2*i], c2 = needle[2*i+1];
		if(c1 == '.') {
			args.reference[i] = 0;
			args.pattern[i] = 0;
		}
		else {
			args.reference[i] = (hexa(c1) << 4) | hexa(c2);
			args.pattern[i] = 1;
		}
	}

	while(!parser.at_end()) {
		parser.accept_options();
		args.regions.push_back(parser.region(session.current_space));
	}

	parser.accept_options();
	parser.end();

	return args;
}

void _afh(Session &session, char const *reference, char const *pattern,
	size_t size, int align, int distance, std::vector<MemoryRegion> &regions)
{
	session.require_vspace();

	/* Default values */
	if(distance < 0) distance = 32;
	if(align <= 0) align = 1;

	VirtualSpace const &v = *session.current_space;

	/* If no region is specified, explore the regions for all bindings */
	if(regions.size() == 0) {
		for(auto &b: v.bindings())
			regions.push_back(b.region);
	}

	int match_count = 0;
	bool output_started = false;

	/* Matches are not shown right away because if they are close enough a
	   single local hexdump will show several of them */
	std::vector<std::pair<uint32_t,int>> pending;

	for(auto const &r: regions) {
		uint32_t region_size = r.end - r.start;
		char const *data = v.translate(r.start, region_size);
		if(!data) throw CommandError("region 0x{:08x} .. 0x{:08x} is not "
			"fully bound", r.start, r.end);

		/* Reach the required alignemnt */
		int i = 0;
		while((r.start + i) % align != 0) i++;

		/* Search patterns for (size) bytes inside (data) */
		for(; i <= (int)region_size - (int)size; i += align) {
			if(!matches(data + i, reference, pattern, size)) continue;

			uint32_t start = r.start + i;

			/* Flush pending matches if this new match is far away */
			if(pending.size() > 0) {
				auto const &p = pending[pending.size() - 1];
				if(p.first + p.second + distance < start) {
					Range r;
					r.start = pending[0].first - distance;
					r.end = p.first + p.second + distance;

					if(output_started) fmt::print("...\n");
					_h_hexdump(session, r, pending);
					output_started = true;
					pending.clear();
				}
			}

			pending.emplace_back(start, size);
			match_count++;
			start += size;

			if(match_count >= 128) break;
		}
		if(match_count >= 128) break;
	}

	/* Print the last pending elements */
	if(pending.size()) {
		auto const &p = pending[pending.size() - 1];
		Range r = { pending[0].first-distance, p.first+p.second+distance };
		if(output_started) fmt::print("...\n");
		_h_hexdump(session, r, pending);
	}

	if(match_count == 0)
		fmt::print("No occurrence found.\n");
	else if(match_count < 128)
		fmt::print("{} occurences found.\n", match_count);
	else
		fmt::print("Stopped after 128 occurrences.\n");
}

//---
// af4
//---

struct _af4_args {
	uint32_t value;
	std::vector<MemoryRegion> regions;
};

static _af4_args parse_af4(Session &session, Parser &parser)
{
	_af4_args args;
	args.value = parser.expr(session.current_space);

	while(!parser.at_end()) {
		args.regions.push_back(parser.region(session.current_space));
	}

	parser.end();
	return args;
}

void _af4(Session &session, uint32_t value, std::vector<MemoryRegion> &regions)
{
	uint32_t value_big_endian = htobe32(value);
	char pattern[4] = { 1, 1, 1, 1 };
	_afh(session, (char *)&value_big_endian, pattern, 4, 4, -1, regions);
}

[[gnu::constructor]] static void _(void)
{
	shell_register_command("af4",
		[](Session &s, Parser &p){
			auto args = parse_af4(s, p);
			_af4(s, args.value, args.regions); },
		[](Session &s, Parser &p){ parse_af4(s, p); });

	shell_register_command("afh",
		[](Session &s, Parser &p) {
			auto args = parse_afh(s, p);
			_afh(s, args.reference.c_str(), args.pattern.c_str(), args.size,
				args.align, args.distance, args.regions); },
		[](Session &s, Parser &p){ parse_afh(s, p); });
}