/* SPDX-License-Identifier: MIT */

#include <algorithm>
#include <deque>
#include <errno.h>
#include <inttypes.h>
#include <limits.h>
#include <new>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <vector>

#include "link/assign.hpp"
#include "link/main.hpp"
#include "link/object.hpp"
#include "link/patch.hpp"
#include "link/sdas_obj.hpp"
#include "link/section.hpp"
#include "link/symbol.hpp"

#include "error.hpp"
#include "helpers.hpp"
#include "linkdefs.hpp"
#include "version.hpp"

static std::deque<std::vector<Symbol>> symbolLists;
static std::vector<std::vector<FileStackNode>> nodes;
static std::deque<Assertion> assertions;

// Helper functions for reading object files

// Internal, DO NOT USE.
// For helper wrapper macros defined below, such as `tryReadlong`
#define tryRead(func, type, errval, vartype, var, file, ...) \
	do { \
		FILE *tmpFile = file; \
		type tmpVal = func(tmpFile); \
		/* TODO: maybe mark the condition as `unlikely`; how to do that portably? */ \
		if (tmpVal == (errval)) { \
			errx(__VA_ARGS__, feof(tmpFile) \
						? "Unexpected end of file" \
						: strerror(errno)); \
		} \
		var = (vartype)tmpVal; \
	} while (0)

/*
 * Reads an unsigned long (32-bit) value from a file.
 * @param file The file to read from. This will read 4 bytes from the file.
 * @return The value read, cast to a int64_t, or -1 on failure.
 */
static int64_t readlong(FILE *file)
{
	uint32_t value = 0;

	// Read the little-endian value byte by byte
	for (uint8_t shift = 0; shift < sizeof(value) * CHAR_BIT; shift += 8) {
		int byte = getc(file);

		if (byte == EOF)
			return INT64_MAX;
		// This must be casted to `unsigned`, not `uint8_t`. Rationale:
		// the type of the shift is the type of `byte` after undergoing
		// integer promotion, which would be `int` if this was casted to
		// `uint8_t`, because int is large enough to hold a byte. This
		// however causes values larger than 127 to be too large when
		// shifted, potentially triggering undefined behavior.
		value |= (unsigned int)byte << shift;
	}
	return value;
}

/*
 * Helper macro for reading longs from a file, and errors out if it fails to.
 * Not as a function to avoid overhead in the general case.
 * @param var The variable to stash the number into
 * @param file The file to read from. Its position will be advanced
 * @param ... A format string and related arguments; note that an extra string
 *            argument is provided, the reason for failure
 */
#define tryReadlong(var, file, ...) \
	tryRead(readlong, int64_t, INT64_MAX, long, var, file, __VA_ARGS__)

// There is no `readbyte`, just use `fgetc` or `getc`.

/*
 * Helper macro for reading bytes from a file, and errors out if it fails to.
 * Not as a function to avoid overhead in the general case.
 * @param var The variable to stash the number into
 * @param file The file to read from. Its position will be advanced
 * @param ... A format string and related arguments; note that an extra string
 *            argument is provided, the reason for failure
 */
#define tryGetc(type, var, file, ...) \
	tryRead(getc, int, EOF, type, var, file, __VA_ARGS__)

/*
 * Helper macro for readings '\0'-terminated strings from a file, and errors out if it fails to.
 * Not as a function to avoid overhead in the general case.
 * @param var The variable to stash the string into
 * @param file The file to read from. Its position will be advanced
 * @param ... A format string and related arguments; note that an extra string
 *            argument is provided, the reason for failure
 */
#define tryReadstring(var, file, ...) \
	do { \
		FILE *tmpFile = file; \
		std::string &tmpVal = var; \
		for (int tmpByte = getc(tmpFile); tmpByte != '\0'; tmpByte = getc(tmpFile)) { \
			if (tmpByte == EOF) { \
				errx(__VA_ARGS__, feof(tmpFile) \
						  ? "Unexpected end of file" \
						  : strerror(errno)); \
			} else { \
				tmpVal.push_back(tmpByte); \
			} \
		}; \
	} while (0)

// Functions to parse object files

/*
 * Reads a file stack node form a file.
 * @param file The file to read from
 * @param nodes The file's array of nodes
 * @param i The ID of the node in the array
 * @param fileName The filename to report in errors
 */
static void readFileStackNode(FILE *file, std::vector<FileStackNode> &fileNodes, uint32_t i,
			      char const *fileName)
{
	FileStackNode &node = fileNodes[i];
	uint32_t parentID;

	tryReadlong(parentID, file,
		    "%s: Cannot read node #%" PRIu32 "'s parent ID: %s", fileName, i);
	node.parent = parentID != (uint32_t)-1 ? &fileNodes[parentID] : NULL;
	tryReadlong(node.lineNo, file,
		    "%s: Cannot read node #%" PRIu32 "'s line number: %s", fileName, i);
	tryGetc(enum FileStackNodeType, node.type, file,
		"%s: Cannot read node #%" PRIu32 "'s type: %s", fileName, i);
	switch (node.type) {
	case NODE_FILE:
	case NODE_MACRO:
		node.data = "";
		tryReadstring(node.name(), file,
			      "%s: Cannot read node #%" PRIu32 "'s file name: %s", fileName, i);
		break;

		uint32_t depth;
	case NODE_REPT:
		tryReadlong(depth, file,
			    "%s: Cannot read node #%" PRIu32 "'s rept depth: %s", fileName, i);
		node.data = std::vector<uint32_t>(depth);
		for (uint32_t k = 0; k < depth; k++)
			tryReadlong(node.iters()[k], file,
				    "%s: Cannot read node #%" PRIu32 "'s iter #%" PRIu32 ": %s",
				    fileName, i, k);
		if (!node.parent)
			fatal(NULL, 0, "%s is not a valid object file: root node (#%"
			      PRIu32 ") may not be REPT", fileName, i);
	}
}

/*
 * Reads a symbol from a file.
 * @param file The file to read from
 * @param symbol The symbol to fill
 * @param fileName The filename to report in errors
 */
static void readSymbol(FILE *file, Symbol *symbol, char const *fileName,
		       std::vector<FileStackNode> const &fileNodes)
{
	tryReadstring(symbol->name, file, "%s: Cannot read symbol name: %s", fileName);
	tryGetc(enum ExportLevel, symbol->type, file, "%s: Cannot read \"%s\"'s type: %s",
		fileName, symbol->name.c_str());
	// If the symbol is defined in this file, read its definition
	if (symbol->type != SYMTYPE_IMPORT) {
		symbol->objFileName = fileName;
		uint32_t nodeID;

		tryReadlong(nodeID, file, "%s: Cannot read \"%s\"'s node ID: %s",
			   fileName, symbol->name.c_str());
		symbol->src = &fileNodes[nodeID];
		tryReadlong(symbol->lineNo, file, "%s: Cannot read \"%s\"'s line number: %s",
			    fileName, symbol->name.c_str());
		tryReadlong(symbol->sectionID, file, "%s: Cannot read \"%s\"'s section ID: %s",
			    fileName, symbol->name.c_str());
		tryReadlong(symbol->offset, file, "%s: Cannot read \"%s\"'s value: %s",
			    fileName, symbol->name.c_str());
	} else {
		symbol->sectionID = -1;
	}
}

/*
 * Reads a patch from a file.
 * @param file The file to read from
 * @param patch The patch to fill
 * @param fileName The filename to report in errors
 * @param i The number of the patch to report in errors
 */
static void readPatch(FILE *file, Patch *patch, char const *fileName, std::string const &sectName,
		      uint32_t i, std::vector<FileStackNode> const &fileNodes)
{
	uint32_t nodeID, rpnSize;
	enum PatchType type;

	tryReadlong(nodeID, file,
		   "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s node ID: %s",
		   fileName, sectName.c_str(), i);
	patch->src = &fileNodes[nodeID];
	tryReadlong(patch->lineNo, file,
		    "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s line number: %s",
		    fileName, sectName.c_str(), i);
	tryReadlong(patch->offset, file,
		    "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s offset: %s",
		    fileName, sectName.c_str(), i);
	tryReadlong(patch->pcSectionID, file,
		    "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s PC offset: %s",
		    fileName, sectName.c_str(), i);
	tryReadlong(patch->pcOffset, file,
		    "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s PC offset: %s",
		    fileName, sectName.c_str(), i);
	tryGetc(enum PatchType, type, file,
		"%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s type: %s",
		fileName, sectName.c_str(), i);
	patch->type = type;
	tryReadlong(rpnSize, file,
		    "%s: Unable to read \"%s\"'s patch #%" PRIu32 "'s RPN size: %s",
		    fileName, sectName.c_str(), i);

	patch->rpnExpression.resize(rpnSize);
	size_t nbElementsRead = fread(patch->rpnExpression.data(), 1, rpnSize, file);

	if (nbElementsRead != rpnSize)
		errx("%s: Cannot read \"%s\"'s patch #%" PRIu32 "'s RPN expression: %s",
		     fileName, sectName.c_str(), i,
		     feof(file) ? "Unexpected end of file" : strerror(errno));
}

/*
 * Sets a patch's pcSection from its pcSectionID.
 * @param patch The patch to fix
 */
static void linkPatchToPCSect(Patch *patch, std::vector<Section *> const &fileSections)
{
	patch->pcSection = patch->pcSectionID != (uint32_t)-1 ? fileSections[patch->pcSectionID]
							      : NULL;
}

/*
 * Reads a section from a file.
 * @param file The file to read from
 * @param section The section to fill
 * @param fileName The filename to report in errors
 */
static void readSection(FILE *file, Section *section, char const *fileName,
			std::vector<FileStackNode> const &fileNodes)
{
	int32_t tmp;
	uint8_t byte;

	tryReadstring(section->name, file, "%s: Cannot read section name: %s", fileName);
	tryReadlong(tmp, file, "%s: Cannot read \"%s\"'s' size: %s", fileName, section->name.c_str());
	if (tmp < 0 || tmp > UINT16_MAX)
		errx("\"%s\"'s section size (%" PRId32 ") is invalid", section->name.c_str(), tmp);
	section->size = tmp;
	section->offset = 0;
	tryGetc(uint8_t, byte, file, "%s: Cannot read \"%s\"'s type: %s", fileName,
		section->name.c_str());
	section->type = (enum SectionType)(byte & 0x3F);
	if (byte >> 7)
		section->modifier = SECTION_UNION;
	else if (byte >> 6)
		section->modifier = SECTION_FRAGMENT;
	else
		section->modifier = SECTION_NORMAL;
	tryReadlong(tmp, file, "%s: Cannot read \"%s\"'s org: %s", fileName, section->name.c_str());
	section->isAddressFixed = tmp >= 0;
	if (tmp > UINT16_MAX) {
		error(NULL, 0, "\"%s\"'s org is too large (%" PRId32 ")", section->name.c_str(), tmp);
		tmp = UINT16_MAX;
	}
	section->org = tmp;
	tryReadlong(tmp, file, "%s: Cannot read \"%s\"'s bank: %s", fileName, section->name.c_str());
	section->isBankFixed = tmp >= 0;
	section->bank = tmp;
	tryGetc(uint8_t, byte, file, "%s: Cannot read \"%s\"'s alignment: %s", fileName,
		section->name.c_str());
	if (byte > 16)
		byte = 16;
	section->isAlignFixed = byte != 0;
	section->alignMask = (1 << byte) - 1;
	tryReadlong(tmp, file, "%s: Cannot read \"%s\"'s alignment offset: %s", fileName,
		    section->name.c_str());
	if (tmp > UINT16_MAX) {
		error(NULL, 0, "\"%s\"'s alignment offset is too large (%" PRId32 ")",
		      section->name.c_str(), tmp);
		tmp = UINT16_MAX;
	}
	section->alignOfs = tmp;

	if (sect_HasData(section->type)) {
		if (section->size) {
			section->data.resize(section->size);
			if (size_t nbRead = fread(section->data.data(), 1, section->size, file);
			    nbRead != section->size)
				errx("%s: Cannot read \"%s\"'s data: %s", fileName, section->name.c_str(),
				     feof(file) ? "Unexpected end of file" : strerror(errno));
		}

		uint32_t nbPatches;

		tryReadlong(nbPatches, file,
			    "%s: Cannot read \"%s\"'s number of patches: %s", fileName,
			    section->name.c_str());

		section->patches.resize(nbPatches);
		for (uint32_t i = 0; i < nbPatches; i++)
			readPatch(file, &section->patches[i], fileName, section->name, i,
				  fileNodes);
	}
}

/*
 * Links a symbol to a section, keeping the section's symbol list sorted.
 * @param symbol The symbol to link
 * @param section The section to link
 */
static void linkSymToSect(Symbol &symbol, Section *section)
{
	uint32_t a = 0, b = section->symbols.size();

	while (a != b) {
		uint32_t c = (a + b) / 2;

		if (section->symbols[c]->offset > symbol.offset)
			b = c;
		else
			a = c + 1;
	}

	section->symbols.insert(section->symbols.begin() + a, &symbol);
}

/*
 * Reads an assertion from a file
 * @param file The file to read from
 * @param assert The assertion to fill
 * @param fileName The filename to report in errors
 */
static void readAssertion(FILE *file, Assertion *assert, char const *fileName, uint32_t i,
			  std::vector<FileStackNode> const &fileNodes)
{
	char assertName[sizeof("Assertion #4294967295")]; // UINT32_MAX

	snprintf(assertName, sizeof(assertName), "Assertion #%" PRIu32, i);

	readPatch(file, &assert->patch, fileName, assertName, 0, fileNodes);
	tryReadstring(assert->message, file, "%s: Cannot read assertion's message: %s", fileName);
}

static Section *getMainSection(Section *section)
{
	if (section->modifier != SECTION_NORMAL)
		section = sect_GetSection(section->name);

	return section;
}

void obj_ReadFile(char const *fileName, unsigned int fileID)
{
	FILE *file;

	if (strcmp(fileName, "-")) {
		file = fopen(fileName, "rb");
	} else {
		fileName = "<stdin>";
		file = fdopen(STDIN_FILENO, "rb"); // `stdin` is in text mode by default
	}
	if (!file)
		err("Failed to open file \"%s\"", fileName);

	// First, check if the object is a RGBDS object or a SDCC one. If the first byte is 'R',
	// we'll assume it's a RGBDS object file, and otherwise, that it's a SDCC object file.
	int c = getc(file);

	ungetc(c, file); // Guaranteed to work
	switch (c) {
	case EOF:
		fatal(NULL, 0, "File \"%s\" is empty!", fileName);

	case 'R':
		break;

	default: // This is (probably) a SDCC object file, defer the rest of detection to it
		// Since SDCC does not provide line info, everything will be reported as coming from the
		// object file. It's better than nothing.
		nodes[fileID].push_back({
			.parent = NULL,
			.lineNo = 0,
			.type = NODE_FILE,
			.data = fileName
		});

		std::vector<Symbol> &fileSymbols = symbolLists.emplace_front();

		sdobj_ReadFile(&nodes[fileID].back(), file, fileSymbols);
		return;
	}

	// Begin by reading the magic bytes
	int matchedElems;

	if (fscanf(file, RGBDS_OBJECT_VERSION_STRING "%n", &matchedElems) == 1
	 && matchedElems != strlen(RGBDS_OBJECT_VERSION_STRING))
		errx("%s: Not a RGBDS object file", fileName);

	verbosePrint("Reading object file %s\n", fileName);

	uint32_t revNum;

	tryReadlong(revNum, file, "%s: Cannot read revision number: %s", fileName);
	if (revNum != RGBDS_OBJECT_REV)
		errx("%s: Unsupported object file for rgblink %s; try rebuilding \"%s\"%s"
			" (expected revision %d, got %d)", fileName, get_package_version_string(),
			fileName, revNum > RGBDS_OBJECT_REV ? " or updating rgblink" : "",
			RGBDS_OBJECT_REV, revNum);

	uint32_t nbNodes;
	uint32_t nbSymbols;
	uint32_t nbSections;

	tryReadlong(nbSymbols, file, "%s: Cannot read number of symbols: %s", fileName);
	tryReadlong(nbSections, file, "%s: Cannot read number of sections: %s", fileName);

	nbSectionsToAssign += nbSections;

	tryReadlong(nbNodes, file, "%s: Cannot read number of nodes: %s", fileName);
	nodes[fileID].resize(nbNodes);
	verbosePrint("Reading %u nodes...\n", nbNodes);
	for (uint32_t i = nbNodes; i--; )
		readFileStackNode(file, nodes[fileID], i, fileName);

	// This file's symbols, kept to link sections to them
	std::vector<Symbol> &fileSymbols = symbolLists.emplace_front(nbSymbols);
	std::vector<uint32_t> nbSymPerSect(nbSections, 0);

	verbosePrint("Reading %" PRIu32 " symbols...\n", nbSymbols);
	for (uint32_t i = 0; i < nbSymbols; i++) {
		// Read symbol
		Symbol &symbol = fileSymbols[i];

		readSymbol(file, &symbol, fileName, nodes[fileID]);

		if (symbol.type == SYMTYPE_EXPORT)
			sym_AddSymbol(&symbol);
		if (symbol.sectionID != -1)
			nbSymPerSect[symbol.sectionID]++;
	}

	// This file's sections, stored in a table to link symbols to them
	std::vector<Section *> fileSections(nbSections, NULL);

	verbosePrint("Reading %" PRIu32 " sections...\n", nbSections);
	for (uint32_t i = 0; i < nbSections; i++) {
		// Read section
		fileSections[i] = new(std::nothrow) Section();
		if (!fileSections[i])
			err("%s: Failed to create new section", fileName);

		fileSections[i]->nextu = NULL;
		readSection(file, fileSections[i], fileName, nodes[fileID]);
		fileSections[i]->fileSymbols = &fileSymbols;
		fileSections[i]->symbols.reserve(nbSymPerSect[i]);

		sect_AddSection(fileSections[i]);
	}

	// Give patches' PC section pointers to their sections
	for (uint32_t i = 0; i < nbSections; i++) {
		if (sect_HasData(fileSections[i]->type)) {
			for (Patch &patch : fileSections[i]->patches)
				linkPatchToPCSect(&patch, fileSections);
		}
	}

	// Give symbols' section pointers to their sections
	for (uint32_t i = 0; i < nbSymbols; i++) {
		int32_t sectionID = fileSymbols[i].sectionID;

		if (sectionID == -1) {
			fileSymbols[i].section = NULL;
		} else {
			Section *section = fileSections[sectionID];

			// Give the section a pointer to the symbol as well
			linkSymToSect(fileSymbols[i], section);

			if (section->modifier != SECTION_NORMAL) {
				if (section->modifier == SECTION_FRAGMENT)
					// Add the fragment's offset to the symbol's
					fileSymbols[i].offset += section->offset;
				section = getMainSection(section);
			}
			fileSymbols[i].section = section;
		}
	}

	uint32_t nbAsserts;

	tryReadlong(nbAsserts, file, "%s: Cannot read number of assertions: %s", fileName);
	verbosePrint("Reading %" PRIu32 " assertions...\n", nbAsserts);
	for (uint32_t i = 0; i < nbAsserts; i++) {
		Assertion &assertion = assertions.emplace_front();

		readAssertion(file, &assertion, fileName, i, nodes[fileID]);
		linkPatchToPCSect(&assertion.patch, fileSections);
		assertion.fileSymbols = &fileSymbols;
	}

	fclose(file);
}

void obj_DoSanityChecks(void)
{
	sect_DoSanityChecks();
}

void obj_CheckAssertions(void)
{
	patch_CheckAssertions(assertions);
}

void obj_Setup(unsigned int nbFiles)
{
	nodes.resize(nbFiles);
}

static void freeSection(Section *section)
{
	for (Section *next; section; section = next) {
		next = section->nextu;
		delete section;
	};
}

void obj_Cleanup(void)
{
	sect_ForEach(freeSection);
}