#include "accumulation.h"

#include <iostream>
#include <filesystem>
#include <fstream>

namespace fs = std::filesystem;

static std::size_t next_global_target_uid = 0;

std::unordered_map<std::size_t, TARGET> targets;
std::unordered_map<std::string, CU> units;

// Two different types with exactly same field structure will not cause problems if simply randomized together
static bool global_target_field_cmp(const TARGET& a, const TARGET& b) {
	if (a.fields.size() != b.fields.size())
		return false;

	auto ita = a.fields.begin();
	auto itb = b.fields.begin();

	for (; ita != a.fields.end() && itb != b.fields.end(); ++ita, ++itb) {
		const FIELD& fa = ita->second;
		const FIELD& fb = itb->second;

		if (fa.offset != fb.offset)
			return false;
		if (fa.size != fb.size)
			return false;
		if (fa.alignment != fb.alignment)
			return false;
		if (fa.flags != fb.flags)
			return false;
	}

	return true;
}

static bool global_target_cmp(const TARGET& a, const TARGET& b) {
	// Note -> Later, use hash of type (generated by pinpoint plugin)

	if (a.name != b.name)
		return false;

	if (a.size != b.size || !global_target_field_cmp(a, b)) {
		std::cerr << "WARNING: Got different definitions of \"" << a.name << "\" -> detached randomization" << std::endl;
		return false;
	}

	return true;
}

static std::size_t accumulate_global_target(TARGET&& target) {
	for (const auto& [guid, gtarget] : targets) {
		if (global_target_cmp(gtarget, target))
			return guid;
	}

	std::size_t guid = next_global_target_uid++;
	targets.emplace(guid, std::move(target));
	return guid;
}

/*
<sourcefile>.spslr:

SPSLR <CU filename> <CU uid symbol>
target <name> <local uid> <size> <field count>
f <offset> <size> <alignment> <flags>
f <offset> <size> <alignment> <flags>
...
ipin <label> <target uid> <field offset> <immediate size>
ipin <label> <target uid> <field offset> <immediate size>
ipin <label> <target uid> <field offset> <immediate size>
...
dpin <symbol> <offset> <level> <target uid>
dpin <symbol> <offset> <level> <target uid>
...
*/
static bool accumulate_file(const fs::path& path) {
	std::ifstream infile(path);
	if (!infile)
		return false;

	std::string hdr_line;
	if (!std::getline(infile, hdr_line))
		return false;

	std::istringstream hdr_iss(hdr_line);

	std::string hdr_magic, hdr_cu_file, hdr_cu_uid;

	if (!(hdr_iss >> hdr_magic) || hdr_magic != "SPSLR")
		return false;

	if (!(hdr_iss >> hdr_cu_file) || !(hdr_iss >> hdr_cu_uid))
		return false;

	std::cout << "Parsing " << path << std::endl;
	std::cout << "    Source file -> " << hdr_cu_file << std::endl;
	std::cout << "    Unit UID    -> " << hdr_cu_uid << std::endl;

	if (units.contains(hdr_cu_uid)) {
		std::cerr << "Duplicate unit UID -> " << hdr_cu_uid << std::endl;
		return false;
	}

	units.emplace(hdr_cu_uid, CU{});
	CU& cu = units.at(hdr_cu_uid);

	std::string line;
	while (std::getline(infile, line)) {
		if (line.empty())
			continue;

		std::istringstream iss(line);
		std::string type;

		if (!(iss >> type))
			return false;

		if (type == "target") {
			TARGET target;

			std::size_t local_uid, field_count;

			if (!(iss >> target.name) || !(iss >> local_uid) || !(iss >> target.size) || !(iss >> field_count)) {
				std::cerr << "Invalid target line" << std::endl;
				return false;
			}

			if (cu.local_targets.contains(local_uid)) {
				std::cerr << "Duplicate local target uid -> " << local_uid << std::endl;
				return false;
			}

			for (std::size_t i = 0; i < field_count; i++) {
				std::string fline;
				if (!std::getline(infile, fline)) {
					std::cerr << "Missing field declaration line!" << std::endl;
					return false;
				}

				std::istringstream fiss(fline);

				std::string ftype;
				if (!(fiss >> ftype) || ftype != "f") {
					std::cerr << "Missing field declaration line, got \"" << ftype << "\" instead!" << std::endl;
					return false;
				}

				FIELD field;
				if (!(fiss >> field.offset) || !(fiss >> field.size) || !(fiss >> field.alignment) || !(fiss >> field.flags)) {
					std::cerr << "Failed to parse field declaration line!" << std::endl;
					return false;
				}

				// Note -> could do sanity checks here

				if (target.fields.contains(field.offset)) {
					std::cerr << "Duplicate field offset!" << std::endl;
					return false;
				}

				target.fields.emplace(field.offset, field);
			}

			auto fit = target.fields.begin();
			for (std::size_t i = 0; i < field_count; i++) {
				fit->second.idx = i;
				fit++;
			}

			std::size_t global_target_uid = accumulate_global_target(std::move(target));
			cu.local_targets.emplace(local_uid, global_target_uid);
			continue;
		} else if (type == "ipin") {
			IPIN ipin;
			if (!(iss >> ipin.symbol) || !(iss >> ipin.local_target) ||
				!(iss >> ipin.field_offset) || !(iss >> ipin.imm_size)) {
				std::cerr << "Failed to parse ipin declaration!" << std::endl;
				return false;
			}

			if (!cu.local_targets.contains(ipin.local_target)) {
				std::cerr << "Ipin references target that has not yet been parsed!" << std::endl;
				return false;
			}

			if (cu.ipins.contains(ipin.symbol)) {
				std::cerr << "Duplicate ipin for symbol \"" << ipin.symbol << "\"!" << std::endl;
				return false;
			}

			cu.ipins.emplace(ipin.symbol, ipin);
			continue;
		} else if (type == "dpin") {
			std::string symbol;
			DPIN::COMPONENT comp;

			if (!(iss >> symbol) || !(iss >> comp.offset) ||
					!(iss >> comp.level) || !(iss >> comp.target)) {
				std::cerr << "Failed to parse dpin declaration!" << std::endl;
				return false;
			}

			if (!cu.dpins.contains(symbol))
				cu.dpins.emplace(symbol, DPIN{ .symbol=symbol });

			DPIN& dpin = cu.dpins.at(symbol);

			if (!cu.local_targets.contains(comp.target)) {
				std::cerr << "Dpin references target that has not yet been parsed!" << std::endl;
				return false;
			}

			dpin.components.push_back(comp);
			continue;
		} else {
			std::cerr << "Invalid spslr file line beginning: \"" << type << "\"" << std::endl;
			return false;
		}
	}

	return true;
}

bool accumulate(const std::vector<std::string>& spslr_files) {
	for (const std::string& spslr_file : spslr_files) {
		fs::path p { spslr_file };

		if (!p.string().ends_with(".spslr")) {
			std::cerr << "The file " << p << " does not appear to be a metadata file!" << std::endl;
			return false;
		}

		if (!fs::exists(p) || !fs::is_regular_file(p)) {
			std::cerr << "Failed to open metadata file " << p << "!" << std::endl;
			return false;
		}

		if (!accumulate_file(p)) {
			std::cerr << "Failed to parse metadata file " << p << "!" << std::endl;
			return false;
		}
	}

	return true;
}