/*
 *  SPL - The SPL Programming Language
 *  Copyright (C) 2004, 2005  Clifford Wolf <clifford@clifford.at>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *  asm.c: The assembler interface for easily creating bytecode files.
 */

#include <sys/types.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <stdarg.h>
#include <assert.h>

#include "spl.h"


struct spl_asm *spl_asm_create(void)
{
	struct spl_asm *as = malloc(sizeof(struct spl_asm));
	as->text = malloc(as->text_size_roof = 4096);
	as->data = malloc(as->data_size_roof = 4096);
	as->text_size = as->data_size = 0;
	as->labels = 0;
	return as;
}

void spl_asm_destroy(struct spl_asm *as)
{
	free(as->text);
	free(as->data);
	free(as);
}

int spl_asm_newdata(struct spl_asm *as, const char *arg)
{
	int data_size = strlen(arg)+1;
	int data_pos = 0;

	while ( data_pos < as->data_size ) {
		if ( !strcmp((char*)(as->data+data_pos), arg) )
			goto found_data_match;
		// just using data_pos++ to find more matches
		// data_pos += strlen(as->data+data_pos)+1;
		data_pos++;
	}

	if ( as->data_size + data_size > as->data_size_roof ) {
		while ( as->data_size + data_size > as->data_size_roof ) as->data_size_roof *= 2;
		as->data = realloc(as->data, as->data_size_roof);
	}

	memcpy(as->data + as->data_size, arg, data_size);
	as->data_size += data_size;

found_data_match:
	return data_pos;
}

int spl_asm_add(struct spl_asm *as, unsigned char op, const char *arg)
{
	int text_size = 1;
	int data_pos = 0;
	int ret = as->text_size;

	if ( op >= 0x20 && op < 0x60 ) {
		if ( op == SPL_OP_PUSHC && !strcmp(arg, "0") ) op = SPL_OP_ZERO;
		if ( op == SPL_OP_PUSHC && !strcmp(arg, "1") ) op = SPL_OP_ONE;
	}

	if ( op < 0x60 ) {
		text_size = 5;
		if (op >= 0x20)
			data_pos = spl_asm_newdata(as, arg);
	}

	if ( as->text_size + text_size > as->text_size_roof ) {
		while ( as->text_size + text_size > as->text_size_roof ) as->text_size_roof *= 2;
		as->text = realloc(as->text, as->text_size_roof);
	}

	*(as->text+as->text_size) = op;
	as->text_size++;

	if ( op < 0x60 ) {
		spl_int_to_bytes(op >= 0x20 ? data_pos : 0, 4, as->text+as->text_size);
		as->text_size += 4;
	}

	return ret;
}

void spl_asm_setaddr(struct spl_asm *as, int pos, int addr)
{
	spl_int_to_bytes(addr - (pos+5), 4, as->text+pos+1);
}

void spl_asm_shuffle(struct spl_asm *as, ...)
{
	struct shuffle_ent;
	struct shuffle_ent {
		struct shuffle_ent *next;
		unsigned char *from, *to, *buffer;
		int size;
	};

	struct shuffle_ent *list = 0;
	va_list ap;
	int arg;

	va_start(ap, as);
	while ( (arg=va_arg(ap, int)) >= 0 ) {
		struct shuffle_ent *e = malloc(sizeof(struct shuffle_ent));

		e->next = list;
		list = e;

		e->size = arg;
		e->from = as->text+va_arg(ap, int);
		e->to = as->text+va_arg(ap, int);

		e->buffer = malloc(e->size);
		memcpy(e->buffer, e->from, e->size);
	}

	while ( list ) {
		struct shuffle_ent *e = list;
		struct spl_asm_label *l = as->labels;

		memcpy(e->to, e->buffer, e->size);

		while (l) {
			if ( (l->position >= e->from - as->text) && (l->position < e->from - as->text + e->size) )
				l->position += e->to - e->from;
			for (int i=0; i<l->refc; i++)
				if ( (l->ref[i] >= e->from - as->text) && (l->ref[i] < e->from - as->text + e->size) )
					l->ref[i] += e->to - e->from;
			l = l->next;
		}

		list = e->next;
		free(e->buffer);
		free(e);
	}
}

static void spl_asm_fixup_text(struct spl_asm *as)
{
	int i;

	for (i=0; i<as->text_size; i++) {
		if ( as->text[i] >= 0x20 && as->text[i] < 0x60 ) {
			int arg_size = 4 - (as->text[i] & 3);
			int arg = spl_bytes_to_int(arg_size, as->text+i+1);
			spl_int_to_bytes(arg + (as->text_size - (i+arg_size+1)), arg_size, as->text+i+1);
			i += arg_size;
		} else
		if ( as->text[i] < 0x20 ) {
			int arg_size = 4 - (as->text[i] & 3);
			i += arg_size;
		}
	}
}

int spl_asm_write(struct spl_asm *as, int fd)
{
	int i, rc;
	int ret = 0;

	spl_asm_fixup_text(as);

	for (i=0; i<16; i+=rc) {
		rc = write(fd, SPL_SIGNATURE+i, 16-i);
		if ( rc <= 0 ) { ret = -1; goto got_error; }
	}

	for (i=0; i<as->text_size; i+=rc) {
		rc = write(fd, as->text+i, as->text_size-i);
		if ( rc <= 0 ) { ret = -1; goto got_error; }
	}

	for (i=0; i<as->data_size; i+=rc) {
		rc = write(fd, as->data+i, as->data_size-i);
		if ( rc <= 0 ) { ret = -1; goto got_error; }
	}

#if 0
	int pagesize = getpagesize();
	int filler = pagesize - (as->text_size + as->data_size) % pagesize;
	if (filler == pagesize) filler = 0;

	if ( filler ) {
		char *cbuf = calloc(1, filler);
		for (i=0; i<filler; i+=rc) {
			rc = write(fd, cbuf+i, filler-i);
			if ( rc <= 0 ) { ret = -1; free(cbuf); goto got_error; }
		}
		free(cbuf);
	}
#endif

got_error:
	as->text_size = 0;
	as->data_size = 0;
	return ret;
}

struct spl_code *spl_asm_dump(struct spl_asm *as)
{
	unsigned char *code = malloc(as->text_size + as->data_size+16);
	struct spl_code *c = spl_code_get(0);

	memcpy(code, SPL_SIGNATURE, 16);

	spl_asm_fixup_text(as);
	memcpy(code+16, as->text, as->text_size);
	memcpy(code+as->text_size+16, as->data, as->data_size);

	c->size = as->text_size + as->data_size + 16;
	c->code_type = SPL_CODE_MALLOCED;
	c->code = code;

	as->text_size = 0;
	as->data_size = 0;
	return c;
}

const char *spl_asm_op2txt(int op)
{
	if (op < 0x60) op = op & ~3;

	for (int i=0; spl_ops[i].name; i++)
		if (spl_ops[i].op == op)
			return spl_ops[i].name;
	return "???";
}

int spl_asm_txt2op(const char *txt)
{
	for (int i=0; spl_ops[i].name; i++)
		if ( !strcasecmp(spl_ops[i].name, txt) )
			return spl_ops[i].op;
	return -1;
}

void spl_asm_print(struct spl_asm *as, int print_data_seg)
{
	printf("# SPL Assembler Dump\n");

	for (int i=0; i<as->text_size; i++) {
		printf(":%-6d %-12s", i+16, spl_asm_op2txt(as->text[i]));
		if (as->text[i] < 0x60) {
			int arg, arg_size = 4 - (as->text[i] & 3);
			arg = spl_bytes_to_int(arg_size, as->text+i+1);
			if (as->text[i] < 0x20) printf(":%-15d", arg + i+arg_size+1 + 16);
			else {
				int oc = 0;
				putchar('"');
				for (int j = 0; (as->data+arg)[j]; j++)
					switch ((as->data+arg)[j])
					{
					case '\\':
						printf("\\\\");
						oc += 2; break;
					case '\"':
						printf("\\\"");
						oc += 2; break;
					case '\n':
						printf("\\n");
						oc += 2; break;
					case '\r':
						printf("\\r");
						oc += 2; break;
					case '\t':
						printf("\\t");
						oc += 2; break;
					default:
						putchar((as->data+arg)[j]);
						oc++;
					}
				putchar('"');
				printf("%*s", oc < 14 ? 14-oc : 0, "");
			}
			printf(" # (%d byte arg) %d", arg_size, arg);
			i += arg_size;
		}
		printf("\n");
	}

	if (print_data_seg) {
		printf("# Data Segment:\n");
		for (int i=0; i<as->data_size; i++) {
			printf("# %5d: \"", i);
			while(as->data[i] && i < as->data_size) {
				switch (as->data[i])
				{
				case '\\':
					printf("\\\\");
					break;
				case '\"':
					printf("\\\"");
					break;
				case '\n':
					printf("\\n");
					break;
				case '\r':
					printf("\\r");
					break;
				case '\t':
					printf("\\t");
					break;
				default:
					putchar(as->data[i]);
				}
				i++;
			}
			printf("\"\n");
		}
	}

	printf("# Total size: %d bytes (%d text, %d data).\n",
		as->text_size + as->data_size, as->text_size, as->data_size);
}

static struct spl_asm_label *get_label(struct spl_asm *as, const char *name)
{
	struct spl_asm_label *l = as->labels;

	while (l) {
		if (!strcmp(l->name, name)) return l;
		l = l->next;
	}

	l = calloc(1, sizeof(struct spl_asm_label));
	l->name = strdup(name);
	l->position = -1;
	l->next = as->labels;
	as->labels = l;
	return l;
}

extern int spl_asm_setlabel(struct spl_asm *as, char *label, int pos)
{
	struct spl_asm_label *l = get_label(as, label);
	if ( l->position >= 0 ) {
		spl_report(SPL_REPORT_ASSEMBLER, as, "Label '%s' redefined!\n", l->name);
		return -1;
	}
	l->position = pos;
	return 0;
}

extern void spl_asm_reflabel(struct spl_asm *as, char *label, int pos)
{
	struct spl_asm_label *l = get_label(as, label);
	l->ref[l->refc++] = pos;
}

extern int spl_asm_resolve_labels(struct spl_asm *as)
{
	struct spl_asm_label *o, *l = as->labels;
	while (l) {
		for (int i=0; i < l->refc; i++) {
			if ( l->position == -1) {
				spl_report(SPL_REPORT_ASSEMBLER, as, "Label '%s' not defined!\n", l->name);
				return -1;
			}
			spl_asm_setaddr(as, l->ref[i], l->position);
		}
		l = (o=l)->next;
		free(o->name);
		free(o);
	}
	as->labels = 0;
	return 0;
}

extern int spl_asm_parse_line(struct spl_asm *as, const char *line)
{
	int this_position = as->text_size;
	line += strspn(line, " \t\n");

	while ( *line == ':' ) {
		int label_len = strcspn(++line, " \t\n");
		char label[label_len+1];

		memcpy(label, line, label_len);
		label[label_len] = 0;

		line += label_len;
		line += strspn(line, " \t\n");

		spl_asm_setlabel(as, label, as->text_size);
	}

	if ( !*line || *line == '#' )
		return 0;

	{
		int op_name_len = strcspn(line, " \t\n");
		char op_name[op_name_len+1];

		memcpy(op_name, line, op_name_len);
		op_name[op_name_len] = 0;

		line += op_name_len;
		line += strspn(line, " \t\n");

		int op = spl_asm_txt2op(op_name);

		if ( op < 0 ) {
			spl_report(SPL_REPORT_ASSEMBLER, as, "Opcode %s not defined!\n", op_name);
			return -1;
		}

		if ( op < 0x20 ) {
			if ( *line == ':' ) {
				int label_len = strcspn(++line, " \t\n");
				char label[label_len+1];

				memcpy(label, line, label_len);
				label[label_len] = 0;

				line += label_len;
				line += strspn(line, " \t\n");

				spl_asm_reflabel(as, label, as->text_size);
				spl_asm_add(as, op, 0);
			} else if ( *line ) {
				int arg_len = strcspn(line, " \t\n");
				char arg[arg_len+1];

				memcpy(arg, line, arg_len);
				arg[arg_len] = 0;

				line += op_name_len;
				line += strspn(line, " \t\n");

				int addr = strtol(arg, 0, 0);
				spl_asm_setaddr(as, spl_asm_add(as, op, 0), addr);
			} else {
				spl_report(SPL_REPORT_ASSEMBLER, as, "Argument for opcode %s missing!\n", op_name);
				return -1;
			}
		} else if ( op < 0x60 ) {
			if ( *line == '"' ) {
				int arg_len = 0;

				line++;
				for (int i=0; line[i] && line[i] != '"'; i++) {
					if (line[i] == '\\' && line[i+1]) i++;
					arg_len++;
				}

				char arg[arg_len+1];

				for (int i=0; *line && *line != '"'; i++, line++) {
					if (*line == '\\' && line[1])
						switch (*(++line))
						{
						case 'n': arg[i] = '\n'; break;
						case 'r': arg[i] = '\r'; break;
						case 't': arg[i] = '\t'; break;
						default:  arg[i] = *line;
						}
					else
						arg[i] = *line;
				}

				arg[arg_len] = 0;

				if (*line == '"') line++;
				line += strspn(line, " \t\n");

				spl_asm_add(as, op, arg);
			} else if ( *line ) {
				int arg_len = strcspn(line, " \t\n");
				char arg[arg_len+1];

				memcpy(arg, line, arg_len);
				arg[arg_len] = 0;

				line += arg_len;
				line += strspn(line, " \t\n");

				spl_asm_add(as, op, arg);
			} else {
				spl_report(SPL_REPORT_ASSEMBLER, as, "Argument for opcode %s missing!\n", op_name);
				return -1;
			}
		} else
			spl_asm_add(as, op, 0);
	}

	if ( *line && *line != '#' ) {
		spl_report(SPL_REPORT_ASSEMBLER, as, "Ignoring cruft (%s) at end of line!\n", line);
		return -1;
	}

	return this_position;
}