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

#include <fxos/disassembly.h>
#include <fxos/vspace.h>
#include <fxos/util/Timer.h>
#include <fxos/util/log.h>
#include <fxos/passes/cfg.h>
#include <fxos/passes/pcrel.h>
#include <fxos/passes/syscall.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("align", [&args](std::string const &value){
		args.align = atoi(value.c_str());
	});
	parser.option("distance", [&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)
{
	if(!session.current_space)
		return;

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

	VirtualSpace &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) {
		char const *data = v.translate(r.start, r.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)r.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 - 1;

					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);
		}
		pending.clear();
	}

	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);
}

//---
// ad
//---

static void ad_disassemble_all(VirtualSpace &space,
	std::vector<uint32_t> const &addresses, bool force)
{
	int successes=0, errors=0;
	Timer timer;

	/* Analyze the CFGs of all functions */

	timer.start();
	CfgPass cfg_pass(space.disasm);

	/* We collect subfunction addresses while running the pass */
	for(int i = 0; i < (int)addresses.size(); i++) {
		uint32_t entry = addresses[i];
		printr("[cfg %d/%zu] Disassembling 0x%08x...",
			i+1, addresses.size(), entry);
		if(!cfg_pass.exploreFunction(entry)) {
			FxOS_log(ERR, "while processing 0x%08x", entry);
			errors++;
			if(!force) return;
		}
		else {
			for(Claim const &c: cfg_pass.resultClaims())
				space.disasm.addExclusiveClaim(c);
			successes++;
		}
	}
	timer.stop();
	printf("\n");
	FxOS_log(LOG, "Finished pass <cfg> in %s", timer.format_time());

	/* Annotate all decoded instructions with pcrel/syscall
	   TODO: analyze only the functions, if possible */
	printr("[pcrel] Resolving PC-relative addressing modes...");
	timer.restart();
	PcrelPass pcrel_pass(space.disasm);
	if(!pcrel_pass.analyzeAllInstructions()) {
		errors++;
		if(!force) return;
	}
	timer.stop();
	printf("\n");
	FxOS_log(LOG, "Finished pass <pcrel> in %s", timer.format_time());

	printr("[syscall] Finding syscall references...");
	timer.restart();
	OS *os = space.os_analysis();
	if(os) {
		SyscallPass syscall_pass(space.disasm, os);
		if(!syscall_pass.analyzeAllInstructions()) {
			errors++;
			if(!force) return;
		}
	}
	timer.stop();
	printf("\n");
	FxOS_log(LOG, "Finished pass <syscall> in %s", timer.format_time());

	printf("Successfully analyzed %d functions (%d errors)\n",
		successes, errors);

	/* TODO: Get subfunction addresses by abstract interpretation and keep
	   going recursively */
}

static std::vector<uint32_t> parse_ad(Session &session, Parser &parser)
{
	if(!session.current_space)
		return std::vector<uint32_t>();

	std::vector<uint32_t> addresses;
	do {
		addresses.push_back(parser.expr(session.current_space));
	}
	while(!parser.at_end());

	parser.end();
	return addresses;
}

void _ad(Session &session, std::vector<uint32_t> const &addresses)
{
	if(!session.current_space)
		return;

	VirtualSpace &space = *session.current_space;
	ad_disassemble_all(space, addresses, false);
}

//--
// ads
//---

static void parse_ads(Session &session, Parser &parser)
{
	if(!session.current_space)
		return;

	parser.end();
}

void _ads(Session &session)
{
	if(!session.current_space)
		return;

	VirtualSpace &space = *session.current_space;
	OS *os = space.os_analysis();

	if(!os) {
		printf("ads: OS analysis failed, cannot enumerate syscalls");
		return;
	}

	std::vector<uint32_t> addresses;
	for(int i = 0; i < os->syscall_count(); i++)
		addresses.push_back(os->syscall(i));

	ad_disassemble_all(space, addresses, true);
}

//---
// Command definitions
//---

static ShellCommand _af4_cmd("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); },
	"Analysis: Find 4-aligned value", R"(
af4 <value> [<regions>...]

Searches mapped memory for a 4-aligned 32-bit value. If regions are specified
(by name or the <start>:<size> syntax), searches these regions. Otherwise,
searches every binding in the current virtual space.

af4 0xb4000000 ROM 0x88000000:512
  Searches occurrences of 0xb4000000, 4-aligned, within mapped ROM and within
  the first 512 bytes of RAM.

af4 0x43415349
  Seacrhes for the 4-aligned string "CASI" in all currently bound regions.
)");

static ShellCommand _afh_cmd("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); },
	"Analysis: Find Hexadecimal pattern", R"(
afh [align=<value>] "<pattern>" [<regions>...]

Searches mapped memory for a hexadecimal pattern with hole bytes. The pattern
should be a hexadecimal string like "0123..6789..cd", consisting either of
explicit hexadecimal values or hole bytes marked by "..". The holes must start
and end at byte boundaries.

An occurrence of the pattern is any sequence of bytes that matches the explicit
bytes exactly, with no constraint on the hole bytes (ie. any value matches).

If regions are specified, searches these regions. Otherwise, searches every
binding in the current virtual space. If align=<align> is specified, only
considers addresses that are multiples of the specified alignment value.

afh "434153494f......00" ROM
  Searches strings ending with "CASIO" followed by up to 3 characters in ROM.

afh align=4 "b4..0000"
  Searches 4-aligned values close to the display interface 0xb4000000 within
  all currently bound regions.
)");

static ShellCommand _ad_cmd("ad",
	[](Session &s, Parser &p) {
		auto addresses = parse_ad(s, p);
		_ad(s, addresses); },
	[](Session &s, Parser &p) { parse_ad(s, p); },
	"Analysis: Disassemble", R"(
ad [<addresses>...]

Disassemble the given set of addresses into the current virtual space's main
disassembly. The main disassembly is used for OS-wide features like cross-
reference search or call graphs.
)");

static ShellCommand _ads_cmd("ads",
	[](Session &s, Parser &p) {
		parse_ads(s, p);
		_ads(s); },
	[](Session &s, Parser &p) { parse_ads(s, p); },
	"Analysis: Disassemble all Syscalls", R"(
ads

Disassembles all syscalls entries using ad, which stores the results in the
current virtual space's main disassembly. Unlike ad, this commands continues
even if some syscalls fail to disassemble.
)");