#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <time.h>

#include "bas.h"

/* remove/comment to turn off debug printing */
#define DEBUG

#ifdef DEBUG
#  define IFDEBUG(x) x
#else
#  define IFDEBUG(x)
#endif

int startlineno = 10;
int increment   = 10;
int limit       = 0;

unsigned short *linerefs[32768];
int linerefcounts[32768];

void print_help(void) {
	fprintf(stderr, "Usage: %s [-v] [-s start-lineno] [-i increment] [-f first-lineno] <inputfile> <outputfile>\n", self);
	fprintf(stderr, "        -v: Verbose.\n");
	fprintf(stderr, "  -s <num>: Starting line number (default: 10).\n");
	fprintf(stderr, "  -i <num>: Increment (default: 10).\n");
	fprintf(stderr, "  -f <num>: Don't renumber lines less than <num> (default: 0).\n");
}

unsigned short getlineno(char opt, const char *arg) {
	int lineno;
	char *e;

	lineno = (int)strtol(arg, &e, 10);

	if(*e) {
		fprintf(stderr, "%s: Invalid line number for -%c option: %s is not a number.\n",
				self, opt, arg);
		exit(1);
	}

	if(lineno < 0 || lineno > 32767) {
		fprintf(stderr, "%s: Invalid line number for -%c option: %d > 32767.\n",
				self, opt, lineno);
		exit(1);
	}

	return ((unsigned short)lineno);
}

void parse_args(int argc, char **argv) {
	int opt;

	while( (opt = getopt(argc, argv, "vs:i:f:")) != -1) {
		switch(opt) {
			case 'v': verbose = 1; break;
			case 's': startlineno = getlineno(opt, optarg); break;
			case 'i': increment   = getlineno(opt, optarg); break;
			case 'f': limit       = getlineno(opt, optarg); break;
			default:
				print_help();
				exit(1);
		}
	}

	if(optind >= argc)
		die("No input file given (use - for stdin).");
	else
		open_input(argv[optind]);

	if(++optind >= argc)
		die("No output file given (use - for stdout).");
	else
		output_filename = argv[optind];
}

unsigned char bcd2int(unsigned char bcd) {
	return (bcd >> 4) * 10 + (bcd & 0x0f);
}

unsigned char int2bcd(unsigned char i) {
	return ((i / 10) << 4) | (i % 10);
}

unsigned short fp2int(const unsigned char *fp) {
	unsigned short result = 0;

	/* examine the exponent/sign byte */
	if(fp[0] == 0) return 0; /* special case */
	if(fp[0] & 0x80) die("negative numbers not supported");

	switch(fp[0]) {
		case 0x40:
			result = bcd2int(fp[1]); break;
		case 0x41:
			result = bcd2int(fp[1]) * 100 + bcd2int(fp[2]); break;
		case 0x42:
			result = bcd2int(fp[1]) * 10000 + bcd2int(fp[2]) * 100 + bcd2int(fp[3]); break;
		default:
			die("number out of range"); break;
	}

	return result;
}

void int2fp(unsigned short num, unsigned char *fp) {
	memset(fp, 0, 6);

	if(num == 0) return;

	if(num >= 10000) {
		fp[0] = 0x42;
		fp[3] = int2bcd(num % 100);
		num /= 100;
		fp[2] = int2bcd(num % 100);
		num /= 100;
		fp[1] = int2bcd(num);
	} else if(num >= 100) {
		fp[0] = 0x41;
		fp[2] = int2bcd(num % 100);
		num /= 100;
		fp[1] = int2bcd(num);
	} else {
		fp[0] = 0x40;
		fp[1] = int2bcd(num);
	}
}

void addlineref(unsigned short refaddr) {
	int target = fp2int(program + refaddr);
	unsigned short *p = linerefs[target];
	int c = linerefcounts[target];

	IFDEBUG(printf("addlineref: target=%d, 0x%04x\n", target, refaddr));

	if(c) {
		p = realloc(p, sizeof(unsigned short) * (c + 1));
	} else {
		p = malloc(sizeof(unsigned short));
	}

	if(!p) die("Out of memory.");

	linerefs[target] = p;
	p[c] = refaddr;
	c++;
	linerefcounts[target] = c;
}

void printlinerefs(void) {
	int i, j;
	printf("linerefs:\n");
	for(i = 0; i < 32768; i++) {
		if(linerefcounts[i]) {
			printf("%d: ", i);
			for(j = 0; j < linerefcounts[i]; j++) {
				printf("%04x ", linerefs[i][j]);
			}
			putchar('\n');
		}
	}
}

void freelinerefs(void) {
	int i;
	for(i = 0; i < 32768; i++) {
		if(linerefcounts[i]) {
			free(linerefs[i]);
		}
	}
}

/* tokens that can take line numbers:
	Commands:
	GOTO   0x0a
	GO TO  0x0b
	GOSUB  0x0c
	TRAP   0x0d
	LIST   0x04 (but don't bother)
	RESTORE 0x23

	Operators:
	GOTO   0x17 (as in, ON (0x1e) GOTO)
	GOSUB  0x18 (ON = 0x1e again)
	THEN   0x1b (but not really!)

	beware: e.g. GOTO 1000+A should not have the 1000 changed.

	numeric constant introduced with 0x0e, followed by 6 BCD bytes.
	string constant 0x0f, length byte, then (length) bytes.
*/

int is_xfer_cmd(unsigned char tok) {
	int ret;
	switch(tok) {
		case CMD_GOTO:
		case CMD_GO_TO:
		case CMD_GOSUB:
		case CMD_TRAP:
		case CMD_LIST:
		case CMD_RESTORE:
			ret = 1; break;
		default:
			ret = 0; break;
	}
	IFDEBUG(printf("is_xfer_cmd(%02x) == %d\n", tok, ret));
	return ret;
}

int skip_op_token(int pos) {
	switch(program[pos]) {
		case OP_EOS:
			return pos + 2; /* skip next-statement offset */
		case OP_EOL:
			return pos + 3; /* skip 2-byte line number */
		case OP_NUMCONST:
			return pos + 7; /* skip 6-byte BCD float */
		case OP_STRCONST:
			return pos + 2 + program[pos + 1]; /* 2nd byte is string len */
		default:
			return pos + 1;
	}
}

/* ON/GOTO and ON/GOSUB can have any number of arguments, separated
   by OP_COMMA. *Normally* these are simple FP constants, since the
   whole point of ON is to avoid computed line numbers... but they're
   allowed to be expressions so we have to check. */
int handle_on_goto(int lineno, int pos) {
	unsigned char tok, nexttok, main_tok;

	IFDEBUG(printf("handle_on_goto(%02x)\n", pos));

	main_tok = program[pos]; /* save this, for use in 'computed' warning */

	pos++; /* skip GOTO/GOSUB token */

	while(1) {
		tok = program[pos];
		if(tok == OP_EOS || tok == OP_EOL)
			break;
		if(tok == OP_COMMA) {
			pos++;
			continue;
		}
		nexttok = program[pos + 7];
		if(tok != OP_NUMCONST || !(nexttok == OP_COMMA || nexttok == OP_EOS || nexttok == OP_EOL)) {
			fprintf(stderr, "Computed line number in ON/%s at line %d.\n",
					(main_tok == OP_GOTO ? "GOTO" : "GOSUB"), lineno);
			break;
		}
		addlineref(pos + 1);
		pos += 7;
	}

	return pos;
}

/* IF/THEN can be followed by a simple line number (not a full expression)
   or a statement offset (*without* an OP_EOS token!) followed by a
   statement.
   The right way to do this would be to track the statement offsets, but
   this works fine. It relies on the fact that line numbers always have
   an exponent byte of 0x40 to 0x42, and the fact that 0x40 to 0x42 are
   not valid command tokens.
 */
int handle_then(int pos) {
	unsigned char tok1 = program[pos + 1];
	unsigned char tok2 = program[pos + 2];
	if(tok1 == OP_NUMCONST && (tok2 >= 0x40 && tok2 <= 0x42)) {
		addlineref(pos + 2);
		return pos + 7;
	} else {
		return 0;
	}
}

void renumber(void) {
	int pos = codestart, nextpos;
	int lineno, offset;
	int newno, i;
	unsigned char tok;
	unsigned char fpnewno[6];

	/* pass 1: find references to line numbers, flag them. */
	while(pos < filelen) {
		lineno = getword(pos);
		if(lineno == 32768) break;
		offset = program[pos + 2];
		IFDEBUG(printf("checking line %d, pos %04x...\n", lineno, pos));
		if(offset < 6)
			die("Can't renumber a code-protected program, unprotect it first.");
		nextpos = pos + offset;

		pos += 4; /* skip line number, line length, 1st statement length */

		/* loop over the statements in this line */
		while(pos < nextpos) {
			/* every statement starts with a command token */
			tok = program[pos];

			if(tok == CMD_REM || tok == CMD_DATA || tok == CMD_ERROR)
				break; /* ignore rest of line */

			if(is_xfer_cmd(tok)) {
				unsigned char nexttok = program[pos + 8];
				if(program[pos + 1] == OP_NUMCONST && (nexttok == OP_EOS || nexttok == OP_EOL)) {
					addlineref(pos + 2);
				} else {
					fprintf(stderr, "Computed line number at line %d.\n", lineno);
				}
			}

			pos++;

			/* rest of statement is expressions/operators */
			while(pos < nextpos) {
				tok = program[pos];

				if(tok == OP_EOL) break;

				if(tok == OP_EOS) {
					pos += 2; /* end statement, skip statement length of next one */
					break;
				}

				if(tok == OP_GOTO || tok == OP_GOSUB) {
					pos = handle_on_goto(lineno, pos);
					continue;
				} else if(tok == OP_THEN) {
					i = handle_then(pos);
					if(i) {
						pos = i;
						continue;
					} else {
						pos += 2; /* skip statement length */
						break;
					}
				}

				IFDEBUG(printf("tok is %02x, pos was %02x before skip_op_token... ", tok, pos));
				pos = skip_op_token(pos);
				IFDEBUG(printf("pos now %02x\n", pos));
			}
		}

		pos = nextpos; /* point to next line */
	}

	IFDEBUG(printlinerefs());

	/* pass 2: renumber the lines, and update the references in other lines */
	newno = startlineno;
	pos = codestart;
	while(pos < filelen) {
		if(newno >= 32768) {
			fprintf(stderr, "New line number %d out of range, renumber failed.", newno);
			exit(1);
		}

		lineno = getword(pos);
		offset = program[pos + 2];

		if(lineno >= limit) {
			if(lineno == 32768) break;

			IFDEBUG(printf("renumbering line %d as %d, %d refs\n", lineno, newno, linerefcounts[lineno]));

			/* update refs to old line number with new one */
			int2fp(newno, fpnewno);
			for(i = 0; i < linerefcounts[lineno]; i++)
				memmove(program + linerefs[lineno][i], fpnewno, 6);

			/* update the actual line number */
			setword(pos, newno);

			newno += increment;
		}
		pos += offset;
	}

	freelinerefs();
}

int main(int argc, char **argv) {
	set_self(*argv);
	parse_general_args(argc, argv, print_help);
	parse_args(argc, argv);

	readfile();
	parse_header();

	renumber();

	open_output(output_filename);
	writefile();

	return 0;
}