/*
 *	Copyright (c) 2024 Scott Duensing, scott@kangaroopunch.com
 *
 *	Permission is hereby granted, free of charge, to any person obtaining a copy
 *	of this software and associated documentation files (the "Software"), to deal
 *	in the Software without restriction, including without limitation the rights
 *	to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 *	copies of the Software, and to permit persons to whom the Software is
 *	furnished to do so, subject to the following conditions:
 *
 *	The above copyright notice and this permission notice shall be included in
 *	all copies or substantial portions of the Software.
 *
 *	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 *	OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 *	SOFTWARE.
 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <dirent.h>

#include "../shared/util.h"


#define BUFFER_SIZE   4096              // If you have a line of code longer than this, you deserve the crash.
#define BLOCK_COUNT   15                // Could be up to 56 blocks, 8-64.  f256.ld needs updated for more.


unsigned char  _currentBank = 0;        // 0 == lower 64k
unsigned char  _totalBanks  = 0;        // Not including lower 64k
char          *_bankName[BLOCK_COUNT];  // Far blocks (8+), shifted down 8.


void findBank(char *name);
void parseCFile(char *filename, char *targetFile, FILE *trampoline, char *trampolineFile, int swapSlot);
void trimEnd(char *string);


void findBank(char *name) {
	int x;

	// Is this the MAIN segment?
	if (strcmp(name, "MAIN") == 0) {
		_currentBank = 0;
		return;
	}

	// Do we know this segment already?
	for (x=0; x<_totalBanks; x++) {
		if (strcmp(name, _bankName[x]) == 0) {
			_currentBank = x + 8;
			return;
		}
	}

	// We have space?
	if (_totalBanks == BLOCK_COUNT) {
		fprintf(stderr, "Too many segments!\n");
		exit(1);
	}

	// Add this new segment!
	_bankName[_totalBanks++] = strdup(name);
	_currentBank = _totalBanks + 7;
}


void parseCFile(char *filename, char *targetFile, FILE *trampoline, char *trampolineFile, int swapSlot) {
	FILE *in;
	FILE *out;
	int   c;
	char  buffer[BUFFER_SIZE];
	char *temp;
	char *start;
	char *b;
	bool  found;
	int   x;
	int   comments     = 0;
	bool  inComment    = false;
	int   brackets     = 0;
	int   crSinceStart = 0;
	int   pos          = 0;
	char *segDef       = "#define SEGMENT_";

	/*
	 * This parser sucks.
	 *
	 *  - It only handles C, not C++.
	 *  - It can only handle single-line function definitions with the opening '{' on the same line as the function.
	 *  - It always generates trampolines, not just when they're needed.
	 *  - Comments after function definitions will probably break it.
	 *  - Comments after SEGMENT defines WILL break it.
	 *
	 *  Someone should fix it.   :-)
	 */

	in = fopen(filename, "rt");
	if (in == NULL) {
		fclose(trampoline);
		fprintf(stderr, "Cannot read %s!\n", filename);
		exit(1);
	}

	out = fopen(targetFile, "wt");
	if (out == NULL) {
		fclose(in);
		fclose(trampoline);
		fprintf(stderr, "Cannot create %s!\n", targetFile);
		free(targetFile);
		exit(1);
	}

	printf("Processing %s -> %s...\n", filename, targetFile);

	// Add trampoline include automatically.
	fprintf(out, "#include \"%s\"\n\n", trampolineFile);

	do {
		// Read next byte from C input file.
		if ((c = fgetc(in)) == EOF) break;

		if (pos > 0) {
			// Look for '//' comments.
			if ((c == '/') && (pos > 0) && (buffer[pos-1] == '/')) inComment = true;

			// Look for '/*' comment start.
			if ((c == '*') && (pos > 0) && (buffer[pos-1] == '/')) comments++;

			// Look for '*/' comment end.
			if ((c == '/') && (pos > 0) && (buffer[pos-1] == '*')) comments--;
		}

		// Count brackets so we know if we're inside a function or not.
		if ((!inComment) && (comments == 0)) {
			if (c == '{') {
				brackets++;
				crSinceStart = 0;
			}
			if (c == '}') brackets--;
		}

		// End of line?
		if ((c == 13) || (c == 10)) {

			inComment = false;
			crSinceStart++;

			// End the line and trim the tail.
			buffer[pos] = 0;
			trimEnd(buffer);

			if ((brackets == 1) && (crSinceStart == 1)) {

				// Is the last character a '{'?  If we're not in a far memory bank, short-circuit this entire mess.
				if ((_currentBank > 0) && (buffer[strlen(buffer) - 1] == '{')) {
					// Remove it and any spaces.
					buffer[strlen(buffer) - 1] = 0;
					trimEnd(buffer);
					// Functions end with a closing parenthesis.
					if (buffer[strlen(buffer) - 1] == ')') {
						// Function.  Annotate it!
						fprintf(out, "__attribute__((noinline, section(\".block%d\")))\n", _currentBank);

						// Scan forward to find the '(' after the function name.
						b = buffer;
						while (*b != '(') b++;
						// Now go backwards and look for a space.
						start = b;
						while (*start != ' ') start--;
						// Are they returning a pointer?
						found = false;
						if (*(start+1) == '*') {
							found = true;
							++start;
							*start++ = 0;
						} else {
							*start++ = 0;
							found = false;
						}
						// Yank the function name out.
						*b = 0;
						temp = strdup(start);
						*b = '(';
						// Write out new function definition.
						fprintf(out, "%s%cFAR%d_%s {\n", buffer, found ? '*' : ' ', _currentBank, start);
						// Create trampoline macro.
						fprintf(trampoline, "#define %s(...)  ({  \\\n"
											"\t\tunsigned char ___mmu = (unsigned char)*(volatile unsigned char *)%#06x;  \\\n"
											"\t\t*(volatile unsigned char *)%#06x = %d;  \\\n"
											"\t\tFAR%d_%s(__VA_ARGS__);  \\\n"
											"\t\t*(volatile unsigned char *)%#06x = ___mmu;  \\\n"
											"\t})\n\n", temp, swapSlot, swapSlot, _currentBank, _currentBank, temp, swapSlot);
						free(temp);
					} else {
						// Not a function.  Write as is.  (Put the '{' back.)
						fprintf(out, "%s {\n", buffer);
					}
				} else {
					// Didn't find '{'.
					fprintf(out, "%s\n", buffer);
				}

			} else {  // Brackets

				// Is this a segment definition?  "#define SEGMENT_<something>"
				found = false;
				if (pos > strlen(segDef)) {
					found = true;
					for (x=0; x<strlen(segDef); x++) {
						if (buffer[x] != segDef[x]) {
							found = false;
							break;
						}
					}
				}
				if (found) {
					// New segment!
					findBank(&buffer[strlen(segDef)]);
				}

				// Write to output in case they depend on the define for some odd reason.
				fprintf(out, "%s\n", buffer);

			} // Brackets.

			// Reset buffer for next line.
			pos = 0;

			// Makes debugging nicer.
			fflush(out);

		} else {  // EOL

			// Add to buffer.
			buffer[pos++] = c;

		}  // EOL

	} while (1);

	fclose(in);
}


void trimEnd(char *string) {
	int x;

	// Trim end of line.
	for (x=strlen(string)-1; x>0; x--) {
		if ((string[x] == ' ') || (string[x] == 9)) {
			string[x] = 0;
		} else {
			break;
		}
	}
}


int main(int argc, char *argv[]) {
	FILE *out;
	int   x;
	int   nearSlot;
	char *targetDir;
	char *sourceDir;
	char *trampolineFile;
	char *linkerFile;
	char *cFile;
	char *targetFile;
	char *thisDir;
	int   sourceDirOffset;
	int   thisOffset;
	DIR  *dir;
	struct dirent *dirent;

	/*
	 *  Command line:
	 *
	 *  ./overlay [nearSlot#] [targetDir] [sourceDir1] ... {sourceDirX}
	 *
	 */

	if (argc < 4) {
		printf("Usage:  %s [nearSlot#] [targetDir] [sourceDir1] ... {sourceDirX}\n", argv[0]);
		return 1;
	}

	// Find near memory slot to swap into.
	nearSlot = atoi(argv[1]);
	if (nearSlot > 7) {
		printf("ERROR!  Only slots 0-7 are in near memory.\n");
		return 1;
	}
	if (nearSlot == 0) {
		printf("WARNING!  Slot 0 contains the zero page, stack, and MMU!\n");
	}
	if ((nearSlot == 6) || (nearSlot == 7)) {
		printf("WARNING!  Slots 6 and 7 contain the microkernel!\n");
	}

	// Does the target directory exist?
	targetDir = strdup(argv[2]);
	utilFixPathSeparators(&targetDir, false);
	if (!utilMkDirP(targetDir, 0777)) {
		printf("ERROR!  Cannot create target directory %s!\n", targetDir);
		return 1;
	}

	// Create trampoline for all C files.
	trampolineFile = utilCreateString("%s%ctrampoline.h", targetDir, utilGetPathSeparator());
	out = fopen(trampolineFile, "wt");
	if (out == NULL) {
		fprintf(stderr, "ERROR!  Cannot create %s!\n", trampolineFile);
		free(trampolineFile);
		free(targetDir);
		return 1;
	}
	fprintf(out, "// THIS FILE IS AUTOMATICALLY GENERATED.  DO NOT EDIT.\n\n");
	fprintf(out, "#ifndef TRAMPOLINE_H\n");
	fprintf(out, "#define TRAMPOLINE_H\n\n");

	// Find common inital path so we can remove it later.
	sourceDir = strdup(argv[3]);
	utilFixPathSeparators(&sourceDir, true);
	sourceDirOffset = strlen(sourceDir) - 1;
	if (argc > 4) {
		for (x=4; x<argc; x++) {
			thisOffset = 0;
			thisDir = strdup(argv[x]);
			utilFixPathSeparators(&thisDir, true);
			for (thisOffset=0; thisOffset < (strlen(sourceDir) < strlen(thisDir) ? strlen(sourceDir) : strlen(thisDir)); thisOffset++) {
				if (sourceDir[thisOffset] != thisDir[thisOffset]) break;
			}
			if (thisOffset < sourceDirOffset) sourceDirOffset = thisOffset - 1;
			free(thisDir);
		}
	}
	free(sourceDir);

	// Do all the C files in a project, one at a time.
	for (x=3; x<argc; x++) {
		sourceDir = strdup(argv[x]);
		utilFixPathSeparators(&sourceDir, true);
		if ((dir = opendir(sourceDir)) == NULL) {
			fprintf(stderr, "ERROR!  Cannot open directory %s!\n", argv[x]);
			free(trampolineFile);
			free(targetDir);
			free(sourceDir);
			return 1;
		}
		for (;;) {
			if ((dirent = readdir(dir)) == NULL) break;
			if ((dirent->d_type == DT_LNK) || (dirent->d_type == DT_REG)) {
				// Is this a C file?
				if ((dirent->d_name[strlen(dirent->d_name) - 2] == '.') && (dirent->d_name[strlen(dirent->d_name) - 1] == 'c')) {
					cFile = utilCreateString("%s%s", sourceDir, dirent->d_name);
					targetFile = utilCreateString("%s%s%s", targetDir, &sourceDir[sourceDirOffset], dirent->d_name);
					thisDir = utilGetUpToLastPathComponent(targetFile);
					utilFixPathSeparators(&thisDir, false);
					if (!utilMkDirP(thisDir, 0777)) {
						printf("ERROR!  Cannot create target directory %s!\n", thisDir);
						free(thisDir);
						free(trampolineFile);
						free(targetDir);
						free(sourceDir);
						free(cFile);
						closedir(dir);
						return 1;
					}
					free(thisDir);
					parseCFile(cFile, targetFile, out, trampolineFile, nearSlot + 8);
					//printf("%s --> %s\n", cFile, targetFile);
					free(cFile);
					free(targetFile);
				}
			}
		}
		closedir(dir);
		free(sourceDir);
	}

	fprintf(out, "#endif // TRAMPOLINE_H\n");
	fclose(out);
	free(trampolineFile);

	// Generate linker script data.
	linkerFile = utilCreateString("%s%coutput.ld", targetDir, utilGetPathSeparator());
	out = fopen(linkerFile, "wt");
	if (out == NULL) {
		fprintf(stderr, "ERROR!  Cannot create %s!\n", linkerFile);
		free(targetDir);
		free(linkerFile);
		return 1;
	}
	for (x=0; x<_totalBanks; x++) {
		fprintf(out, "  SHORT(%d*0x2000)\n", x+8);
		fprintf(out, "  BYTE(%d/8)\n", x+8);
		fprintf(out, "  SHORT(end_block%d - __block%d_lma)\n", x+8, x+8);
		fprintf(out, "  BYTE(0x00)\n");
		fprintf(out, "  TRIM(block%d)\n\n", x+8);
	}
	fclose(out);
	free(linkerFile);

	// Clean up.
	for (x=0; x<_totalBanks; x++) {
		free(_bankName[x]);
	}
	free(targetDir);

	return 0;
}