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

#include <LIEF/LIEF.hpp>
#include <LIEF/ELF/Builder.hpp>

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

#include <spslr_program.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)!
*/

using namespace LIEF::ELF;

static bool disassemble_ipin(const Section* text, IPIN::HIT& pin);
static bool assemble_patcher_program(uint64_t vaddr_pivot, std::vector<uint8_t>& program);

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<Binary> bin = 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 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;
			}
		}
	}

	// Find __spslr_program as anchor of relative addressing in patcher program
	if (!global_syms.symbols.contains("__spslr_program")) {
		std::cerr << "Unable to locate \"__spslr_program\"!" << std::endl;
		return 1;
	}

	uint64_t spslr_program_ptr_address = global_syms.symbols.at("__spslr_program");

	// Serialize patcher program
	std::vector<uint8_t> patcher_program;
	if (!assemble_patcher_program(spslr_program_ptr_address, patcher_program)) {
		std::cerr << "Failed to assemble patcher program!" << std::endl;
		return 1;
	}

	std::cout << "Generated patcher program of 0x" << std::hex << patcher_program.size()
		<< std::dec << " bytes!" << std::endl;

	// Add new section ".spslr" with patcher program

	uint64_t page_align = 0x1000;
	for (const auto& seg : bin->segments()) {
		if (seg.type() == Segment::TYPE::LOAD) {
			if (seg.alignment() > page_align) page_align = seg.alignment();
		}
	}

	uint64_t max_end_vaddr = 0;
	for (const auto& seg : bin->segments()) {
		if (seg.type() == Segment::TYPE::LOAD) {
			uint64_t endv = seg.virtual_address() + seg.virtual_size();
			if (endv > max_end_vaddr) max_end_vaddr = endv;
		}
	}

	uint64_t new_vaddr = max_end_vaddr;
	if (new_vaddr % page_align != 0)
		new_vaddr += page_align - (new_vaddr % page_align);

	std::cout << "Adding patcher program at 0x" << std::hex << new_vaddr << std::dec << std::endl;

	Segment new_seg;
	new_seg.type(Segment::TYPE::LOAD);
	new_seg.flags(Segment::FLAGS::R); //(ELF_SEGMENT_FLAGS::PF_R);
	new_seg.alignment(page_align);
	new_seg.file_offset(0); // Writer decides where to put it
	new_seg.virtual_address(new_vaddr);
	new_seg.content(patcher_program);
	new_seg.virtual_size(patcher_program.size());

	bin->add(new_seg);

	// Set __spslr_program to (new_vaddr - &__spslr_program)
	for (Segment& seg : bin->segments()) {
		uint64_t start = seg.virtual_address();
		uint64_t end_mem = start + seg.virtual_size();

		if (spslr_program_ptr_address >= start && spslr_program_ptr_address < end_mem) {
			uint64_t offset_within_seg = spslr_program_ptr_address - start;
			uint64_t min_needed = offset_within_seg + sizeof(uint64_t);

			if (min_needed > seg.physical_size())
				seg.physical_size(min_needed);

			break;
		}
	}

	uint64_t relative_new_vaddr = new_vaddr - spslr_program_ptr_address;
	std::vector<uint8_t> program_ptr_bytes;
	program_ptr_bytes.resize(8);
	std::memcpy(program_ptr_bytes.data(), &relative_new_vaddr, 8);
	bin->patch_address(spslr_program_ptr_address, program_ptr_bytes);

	// Output final program
	Builder builder{ *bin };
	builder.build();
	builder.write(out_file);

	std::cout << "Wrote final binary to \"" << out_file << "\"" << std::endl;

	return 0;
}

bool assemble_patcher_program(uint64_t vaddr_pivot, std::vector<uint8_t>& program) {
	uint8_t buf[sizeof(SPSLR_INST)];

	auto append_inst = [&](const SPSLR_INST& inst) -> bool {
		int n = spslr_inst_dump(&inst, buf);
		if (n < 0) {
			std::cerr << "Failed to dump patcher instruction!" << std::endl;
			return false;
		}

		std::size_t program_ptr = program.size();

		program.resize(program.size() + n);
		std::memcpy(program.data() + program_ptr, buf, n);
		return true;
	};

	SPSLR_INST inst;

	// Dump all targets, their fields and command randomization
	for (const auto& [tuid, target] : targets) {
		inst.opcode = SPSLR_TARGET;
		inst.op0.target_uid = tuid;
		inst.op1.target_size = target.size;
		inst.op2.target_fieldcnt = target.fields.size();

		if (!append_inst(inst))
			return false;

		for (const auto& [foff, field] : target.fields) {
			inst.opcode = SPSLR_FIELD;
			inst.op0.field_offset = field.offset;
			inst.op1.field_size = field.size;
			inst.op2.field_alignment = field.alignment;
			inst.op3.field_flags = field.flags; // Note -> This should be done explicitely with SPSLR_FLAG_FIELD_FIXED

			if (!append_inst(inst))
				return false;
		}

		inst.opcode = SPSLR_RANDOMIZE;
		inst.op0.randomize_target = tuid;

		if (!append_inst(inst))
			return false;
	}

	// Dump ipins with addresses relative to vaddr_pivot
	auto append_ipin = [&](const CU& cu, const IPIN& ipin) -> bool {
		inst.opcode = SPSLR_IPATCH;
		inst.op0.ipatch_ptr = ipin.hit->vaddr + ipin.hit->imm_offset - vaddr_pivot;
		inst.op1.ipatch_size = ipin.hit->imm_size;
		inst.op2.ipatch_target = cu.local_targets.at(ipin.local_target);

		if (!targets.at(inst.op2.ipatch_target).fields.contains(ipin.field_offset))
			return false;

		inst.op3.ipatch_field = targets.at(inst.op2.ipatch_target).fields.at(ipin.field_offset).idx;
		return append_inst(inst);
	};

	// Dump dpins with addresses relative to vaddr_pivot and components ordered by level
	auto append_dpin = [&](const DPIN& dpin) -> bool {
		std::list<DPIN::COMPONENT> sorted_components = dpin.components;
		sorted_components.sort([](const DPIN::COMPONENT& a, const DPIN::COMPONENT& b) {
			return a.level > b.level;  // Descending order
		});

		for (const DPIN::COMPONENT& comp : sorted_components) {
			inst.opcode = SPSLR_DPATCH;
			inst.op0.dpatch_ptr = dpin.hit->vaddr + comp.offset - vaddr_pivot;
			inst.op1.dpatch_target = comp.target;

			if (!append_inst(inst))
				return false;
		}

		return true;
	};

	// Dump all local pins
	for (const auto& [cu_uid, cu] : units) {
		for (const auto& [ipin_sym, ipin] : cu.ipins) {
			if (!append_ipin(cu, ipin))
				return false;
		}

		for (const auto& [dpin_sym, dpin] : cu.dpins) {
			DPIN global_dpin;
			for (DPIN::COMPONENT& comp : global_dpin.components)
				comp.target = cu.local_targets.at(comp.target);

			if (!append_dpin(global_dpin))
				return false;
		}
	}

	// Dump global dpins
	for (const auto& [dpin_sym, dpin] : global_dpins) {
		if (!append_dpin(dpin))
			return false;
	}

	// Exit patcher program
	inst.opcode = SPSLR_EXIT;
	if (!append_inst(inst))
		return false;

	return true;
}

bool disassemble_ipin(const 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;
}