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

#include "bas.h"

int protect_vars = 1;
int protect_code = 1;
int shrinktable  = 0;
int varname_char = 0x9b;

/* 32767 END */
unsigned char badcode[] = {
	0xff, 0x7f,   /* line number 32767      */
	0x00,         /* *bad* next-line offset */
	0x06,         /* next-statement offset  */
	0x15,         /* END token              */
	0x16,         /* end-of-line token      */
};

void scramble_vars(void) {
	int i;

	if(!vntable_ok()) {
		fprintf(stderr, "%s: Program already was variable-protected.\n", self);
		exit(2);
	}

	if(shrinktable) {
		if(verbose) fprintf(stderr, "Shrinking variable name table.\n");
		adjust_vntable_size((vvstart - 1) - vnstart, (codestart - vvstart) / 8);
	}

	if(varname_char == -1) srand(time(NULL));

	for(i = vnstart; i < vvstart - 1; i++)
		if(varname_char == -1)
			program[i] = (rand() >> 8) & 0xff;
		else
			program[i] = varname_char & 0xff;

	if(verbose) {
		i -= vnstart;
		if(i) {
			fprintf(stderr, "Replaced %d byte variable name table with ", i);
			if(varname_char == -1)
				fprintf(stderr, "random characters.\n");
			else
				fprintf(stderr, "character $%02x.\n", varname_char);
		} else {
			fprintf(stderr, "Can't protect variables because there are no variables.\n");
		}
	}
}

CALLBACK(bad_offset) {
	fprintf(stderr, "%s: program already was code-protected.\n", self);
	exit(2);
}

unsigned short last_pos = 0;
int last_lineno = -1;

CALLBACK(save_linepos) {
	last_pos = pos;
	if(lineno == 32768) return;
	last_lineno = lineno;
}

/* iterate over all the lines, insert a poisoned line 32767 just
	before line 32768 */
void breakcode(void) {
	int offset;

	on_start_line = save_linepos;
	on_bad_line_length = bad_offset;
	walk_all_code();

	if(last_lineno == -1) die("Can't protect code because there are no lines of code.");
	if(last_lineno == 32767) die("Can't protect code because there is already a line 32767.");

	/* last_pos is now the start of line 32768, move it up to make room for
		the new line */
	offset = sizeof(badcode);
	memmove(program + last_pos + offset, program + last_pos, filelen);

	/* insert new line */
	memmove(program + last_pos, badcode, offset);

	if(verbose)
		fprintf(stderr, "Inserted line 32767 with invalid offset at file offset $%04x.\n", last_pos);

	/* update pointers that would be affected by the code move */
	stmcur += offset;
	starp += offset;
	filelen += offset;
	update_header();
	parse_header();
}

void print_help(void) {
	printf("Usage: %s [-v] [-nc|-nv] [-s] [-x[r|NN]] <inputfile> <outputfile>\n", self);
	printf("  -v: Verbose.\n");
	printf(" -nc: Don't protect code.\n");
	printf(" -nv: Don't protect variable names.\n");
	printf("  -s: Shrink variable name table to min size.\n");
	printf("-xNN: Hex code NN for variable names.\n");
	printf(" -xr: Random variable names.\n");
	printf("Use - as a filename to read from stdin and/or write to stdout.\n");
}

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

	while( (opt = getopt(argc, argv, "vn:x:s")) != -1) {
		switch(opt) {
			case 'v': verbose = 1; break;
			case 's': shrinktable = 1; break;
			case 'n':
				switch(optarg[0]) {
					case 'c': protect_code = 0; break;
					case 'v': protect_vars = 0; break;
					default:
						die("Invalid argument for -n (must be 'c' or 'v').");
				}
				break;
			case 'x':
				xopt_used = 1;
				switch(optarg[0]) {
					case 'r':
						varname_char = -1; break;
					case 0:
						die("-x option requires a hex number or 'r'."); break;
					default:
						{
							char *e;
							varname_char = (int)strtol(optarg, &e, 16);
							if(*e != 0 || varname_char > 0xff)
								fprintf(stderr, "%s: Invalid hex value '%s' for -x option (range is 0 to ff).\n", self, optarg);
						}
				}
				break;
			default:
				print_help();
				exit(1);
		}
	}

	if(!protect_code && !protect_vars) {
		die("Nothing to do: -nc and -nv both given.");
	}

	if(!protect_vars) {
		if(xopt_used)
			die("-x option not valid with -nv.");
		if(shrinktable)
			die("-s option not valid with -nv.");
	}

	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];
}

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

	if(verbose) {
		fprintf(stderr, "Protecting program, ");
		if(protect_vars && !protect_code)
			fprintf(stderr, "variables only.\n");
		else if(protect_code && !protect_vars)
			fprintf(stderr, "code only.\n");
		else
			fprintf(stderr, "both code and variables.\n");
	}
	if(protect_vars) scramble_vars();
	if(protect_code) breakcode();

	open_output(output_filename);
	writefile();
	return 0; /* TODO: meaningful return status */
}
