#include <iostream>
#include <string>
#include <getopt.h>

#include <LIEF/LIEF.hpp>

#include "accumulation.h"
#include "symbol_collection.h"

/*
Notes:
	Datapins for same var/symbol are randomized in order of their level, from bottom of nest to top
	The CU uid symbol helps differentiating between e.g. "file.c" and "sub/file.c" (symbtab has no idea)
	Between CUs, types with the same name HAVE TO HAVE the same layout -> randomized together
	To begin with, anonymous types are not allowed for randomization (later solved with hash(type) instead of name)!
*/

static bool disassemble_ipin(const LIEF::ELF::Section* text, IPIN::HIT& pin);

int main(int argc, char** argv) {
	static option long_options[] = {
		{ "help", no_argument, 0, 0 },
		{ "spslr", required_argument, 0, 0 },
		{ "bin", required_argument, 0, 0 },
		{ "out", required_argument, 0, 0 },
		{ "strip", no_argument, 0, 0 },
		{ 0, 0, 0, 0 }
	};

	int option_index = 0;
	int c;

	std::string spslr_dir, bin_file, out_file;

	while ((c = getopt_long(argc, argv, "", long_options, &option_index)) == 0) {
		const option& opt = long_options[option_index];
		std::string optname { opt.name };

		if (optname == "help") {
			std::cout << "To use spslr_finalize, supply these 3 arguments:" << std::endl;
			std::cout << "    --spslr=<dir> (the directory of .spslr files produced by spslr_pinpoint)" << std::endl;
			std::cout << "    --bin=<file> (the binary compiled with spslr_pinpoint to be finalized)" << std::endl;
			std::cout << "    --out=<file> (the finalized binary file to be written)" << std::endl;
			return 0;
		} else if (optname == "spslr") {
			spslr_dir = std::string{ optarg };
		} else if (optname == "bin") {
			bin_file = std::string{ optarg };
		} else if (optname == "out") {
			out_file = std::string{ optarg };
		} else if (optname == "strip") {
			std::cerr << "Symbol stripping (--strip) is not yet implemented!" << std::endl;
			return 1;
		} else {
			std::cerr << "Invalid option, try \"--help\"!" << std::endl;
			return 1;
		}
	}

	if (spslr_dir.empty()) {
		std::cerr << "Missing spslr directory, supply it via --spslr=<dir>!" << std::endl;
		return 1;
	}

	if (bin_file.empty()) {
		std::cerr << "Missing input file path, supply it via --bin=<file>!" << std::endl;
		return 1;
	}

	if (out_file.empty()) {
		std::cerr << "Missing output file path, supply it via --out=<file>!" << std::endl;
		return 1;
	}

	if (!accumulate(spslr_dir)) {
		std::cerr << "Failed to accumulate data from spslr directory!" << std::endl;
		return 1;
	}

	std::cout << "Gathered a total of " << targets.size() << " distinct targets from "
		<< units.size() << " compilation units!" << std::endl;

	std::unique_ptr<LIEF::ELF::Binary> bin = LIEF::ELF::Parser::parse(bin_file);
	if (!bin) {
		std::cerr << "Failed to parse binary \"" << bin_file << "\"!" << std::endl;
		return 1;
	}

	if (!collect_symbols(bin)) {
		std::cerr << "Failed to collect symbols from \"" << bin_file << "\"!" << std::endl;
		return 1;
	}

	if (!associate_symbols()) {
		std::cerr << "Failed to associate symbols with accumulated data!" << std::endl;
		return 1;
	}

	// For each ipin, disassemble instruction and find immediate offset
	const LIEF::ELF::Section* text = bin->get_section(".text");
	if (!text) {
		std::cerr << "Unable to locate .text section for ipin disassembly!" << std::endl;
		return 1;
	}

	for (auto& [cu_uid, cu] : units) {
		for (auto& [_, ipin] : cu.ipins) {
			if (!ipin.hit.has_value()) {
				std::cerr << "Encountered ipin without vaddr!" << std::endl;
				return 1;
			}

			if (!disassemble_ipin(text, ipin.hit.value())) {
				std::cerr << "Failed to disassemble ipin!" << std::endl;
				return 1;
			}
		}
	}

	// TODO
	/*
	1. Find __spslr_program
	2. Serialize entire patcher program
		-> all pin addresses relative to &__spslr_program
		-> dpatch entry components ordered by level
	3. Add new section .spslr with patcher program
	4. Set __spslr_program to (&.spslr - &__spslr_program)
	5. Output final program
	*/
}

bool disassemble_ipin(const LIEF::ELF::Section* text, IPIN::HIT& pin) {
	if (!text)
		return false;

	uint64_t text_begin = text->virtual_address();
	uint64_t text_size = text->size();
	auto text_data = text->content();

	uint64_t pin_addr = pin.vaddr;

	if (pin_addr < text_begin || pin_addr >= text_begin + text_size)
		return false;

	uint64_t pin_offset = pin_addr - text_begin;

	// 32 bit mov of immediate to 64 bit register: 0x48 0xc7 [8 bit reg] [32 bit immediate]

	if (text_data[pin_offset] != 0x48 || text_data[pin_offset + 1] != 0xc7) {
		std::cerr << "Ipin uses not yet handled instruction!" << std::endl;
		return false;
	}

	pin.imm_offset = 3;
	pin.imm_size = 4;
	return true;
}