#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <assert.h>

#define DBGTIME(tag,block) \
	do { \
		clock_t start=clock(); \
		DBGTIMESTART(start,tag,block); \
	} while(0)

#define DBGTIMESTART(start,tag,block) \
	do { \
		if(params.verbose){ \
			fprintf(stderr,"%s",(tag)); \
		} \
		{block} \
		if(params.verbose){ \
			fprintf(stderr,"  (%lfs)\n",(double)(clock()-start)/CLOCKS_PER_SEC); \
		} \
	} while(0)


__attribute__((noreturn)) void outofmem(void){
	fprintf(stderr,"Memory allocation error!\n");
	exit(1);
}

int uniqid(void){
	static int i=0;
	return i++;
}

void usage(const char *argv0){
	fprintf(stderr,
		"Usage: %s [-h] [-m <memsize>] <srcfile> <dstfile>\n"
		"Reads a Brainfuck program from <srcfile> and saves the compiled executable\n"
		"in <dstfile>.\n"
		"Uses nasm for assembling and gcc for linking.\n"
		"  -h             Show help\n"
		"  -m <memsize>   Specify amount of tape cells (bytes) for the BF code\n"
		"  -v             Show more info about the compile process\n",
		argv0);
}

__attribute__((noreturn)) void usage1(const char *argv0){
	usage(argv0);
	exit(1);
}

typedef struct Params{
	int memsize;
	const char *srcfname,*dstfname;
	bool verbose;
} Params;
Params params={30000,NULL,NULL,false};

void parseargs(int argc,char **argv){
	if(argc<2)usage1(argv[0]);

	for(int i=1;i<argc;i++){
		bool skipnext=false;
		if(argv[i][0]=='-'){
			for(int j=1;argv[i][j];j++){
				switch(argv[i][j]){
					case 'h':
						usage(argv[0]);
						exit(0);

					case 'm':{
						if(i==argc-1)usage1(argv[0]);
						char *endp;
						params.memsize=strtol(argv[i+1],&endp,10);
						if(!argv[i+1][0]||*endp||params.memsize<1)usage1(argv[0]);
						skipnext=true;
						break;
					}

					case 'v':
						params.verbose=true;
						break;
				}
			}
		} else if(params.srcfname==NULL)params.srcfname=argv[i];
		else if(params.dstfname==NULL)params.dstfname=argv[i];
		else usage1(argv[0]);
		if(skipnext)i++;
	}

	if(!params.srcfname||!params.dstfname)usage1(argv[0]);
}

int runcmd(const char *cmd/*,int in,int out*/){
	clock_t start=clock();
	pid_t pid=fork();
	if(pid<0){
		fprintf(stderr,"Cannot fork!\n");
		exit(1);
	}
	if(pid==0){
		/*dup2(in,0);
		dup2(out,1);*/
		execl("/bin/sh","/bin/sh","-c",cmd,NULL);
		__builtin_unreachable();
	} else {
		int statloc;
		DBGTIMESTART(start,cmd,{
			while(true){
				waitpid(pid,&statloc,0);
				if(WIFEXITED(statloc))break;
			}
		});
		return WEXITSTATUS(statloc);
	}
}

FILE* gettempfilew(char name[20]){
	memcpy(name,"/tmp/tmp.XXXXXXXXXX",20);
	int fd=mkstemp(name);
	if(fd<0){
		fprintf(stderr,"Cannot create temp file!\n");
		exit(1);
	}
	FILE *f=fdopen(fd,"w");
	if(!f)outofmem();
	return f;
}

int compilechain(const char *asmfname,const char *dstfname){
	char ofname[20];
	FILE *f=gettempfilew(ofname);
	assert(strlen(ofname)==19);
	fclose(f);
	char buf[25+19+1+strlen(asmfname)+1];
	sprintf(buf,"nasm -f macho64 -w+all -o %s %s",ofname,asmfname);
	int ret=runcmd(buf);
	if(ret)return ret;
	char buf2[4+19+4+strlen(dstfname)+1];
	sprintf(buf2,"gcc %s -o %s",ofname,dstfname);
	ret=runcmd(buf2);
	unlink(ofname);
	if(params.verbose)fprintf(stderr,"unlink(%s)  (ofile)\n",ofname);
	return ret;
}

void writeprologue(FILE *asmf){
	fprintf(asmf,
		"global _main\n"
		"\n"
		"extern _malloc\n"
		"extern _free\n"
		"extern _printf\n"
		"extern _getchar\n"
		"extern _putchar\n"
		"extern _write\n"
		"extern ___stdoutp\n"
		"extern _fflush\n"
		"\n"
		"default rel\n"
		"\n"
		"section .text\n"
		"\n"
		"_main:\n"
		"\tpush rbp  ; prologue\n"
		"\tmov rbp, rsp\n"
		"\tadd rsp, 8\n"
		"\tpush rbx\n"
		"\tmov edi, %d\n"
		"\tcall _malloc\n"
		"\tcmp rax, 0\n"
		"\tjz nomem\n"
		"\tmov rbx, rax ;bufp\n"
		"\tmov [buf], rax ;buf\n"
		"\n",
		params.memsize);
}

void writeprogram(char *source,FILE *asmf){
	int addc=0,shiftc=0;
	//int lastshiftc=0;
	struct stackitem {int id; struct stackitem *next;} *loopstack=NULL;
	int i=0;
	do {
		char c=source[i];
		if(addc&&c!='+'&&c!='-'){
			addc=(addc%256)+256%256;
			if(addc==1)fprintf(asmf,"\tinc byte [rbx]  ; 1 +\n");
			else if(addc==-1)fprintf(asmf,"\tdec byte [rbx]  ; 1 -\n");
			else fprintf(asmf,"\tadd byte [rbx], %d  ; %d %c\n",addc,addc<0?-addc:addc,addc<0?'-':'+');
			addc=0;
		}
		if(shiftc&&c!='>'&&c!='<'){
			shiftc=(shiftc%256)+256%256;
			/*if(lastshiftc>0&&shiftc<0){
				int id=uniqid();
				fprintf(asmf,
					"\tcmp [maxp], rbx\n"
					"\tjge skip%d\n"
					"\tmov [maxp], rbx\n"
					"skip%d:\n",
					id,id);
			}
			lastshiftc=shiftc;*/
			if(shiftc==1)fprintf(asmf,"\tinc rbx  ; 1 >\n");
			else if(shiftc==-1)fprintf(asmf,"\tdec rbx  ; 1 <\n");
			else fprintf(asmf,"\tadd rbx, %d  ; %d %c\n",shiftc,shiftc<0?-shiftc:shiftc,shiftc<0?'<':'>');
			shiftc=0;
		}
		switch(c){
			case '+': addc++; break;
			case '-': addc--; break;
			case '>': shiftc++; break;
			case '<': shiftc--; break;

			case '[':{
				struct stackitem *ls=malloc(sizeof(struct stackitem));
				if(!ls)outofmem();
				ls->id=uniqid();
				ls->next=loopstack;
				loopstack=ls;
				fprintf(asmf,
					"loop%d:  ; [\n"
					"\tcmp byte [rbx], 0\n"
					"\tjz loop%d_end\n"
					"loop%d_body:\n",
					ls->id,ls->id,ls->id);
				break;
			}

			case ']':{
				if(!loopstack){
					fprintf(stderr,"Excess ']' in source\n");
					exit(1);
				}
				fprintf(asmf,
					"\tcmp byte [rbx], 0  ; ]\n"
					"\tjnz loop%d_body\n"
					"loop%d_end:\n",
					loopstack->id,loopstack->id);
				struct stackitem *ls=loopstack->next;
				free(loopstack);
				loopstack=ls;
				break;
			}

			case '.':
				fprintf(asmf,
					"\tmov edi, 0  ; .\n"
					"\tmov dil, byte [rbx]\n"
					"\tcall _putchar\n"
					"\txor edi, edi\n"
					"\tcall _fflush\n");
				break;

			case ',':
				fprintf(asmf,
					"\tcall _getchar  ; ,\n"
					"\tmov byte [rbx], al\n");
				break;

			case '0':
				fprintf(asmf,"\tmov byte [rbx], 0\n");
				break;

			default:
				if(c)assert(false); //should've been stripped
		}
	} while(source[i++]);
	if(loopstack){
		fprintf(stderr,"Excess '[' in source\n");
		exit(1);
	}
}

void writeepilogue(FILE *asmf){
	fprintf(asmf,
		"\n"
		"\tmov rdi, [buf]  ; epilogue\n"
		"\tcall _free\n"
		"\n"
		/*"\tlea rdi, [formatstr] ;output last cell reached\n"
		"\tmov rsi, [maxp]\n"
		"\tsub rsi, [buf]\n"
		"\txor eax, eax\n"
		"\tcall _printf\n"
		"\n"*/
		"\txor eax, eax\n"
		"\tjmp mainend\n"
		"\n"
		"nomem:\n"
		"\tmov edi, 2\n"
		"\tlea rsi, [nomemstr]\n"
		"\tmov rdx, nomemstr.len\n"
		"\tcall _write\n"
		"\tmov eax, 1\n"
		"\n"
		"mainend:\n"
		"\tpop rbx\n"
		"\tmov rsp, rbp\n"
		"\tpop rbp\n"
		"\tret\n"
		"\n"
		"\n"
		"section .data\n"
		"sectalign 8\n"
		"\n"
		"nomemstr: db \"Out of memory!\"\n"
		".len: equ $-nomemstr\n"
		/*"\n"
		"formatstr: db \"Last cell reached: %%d\", 10, 0\n"
		"\n"
		"align 8, db 0\n"
		"maxp: dq 0\n"*/
		"\n"
		"\n"
		"section .bss\n"
		"sectalign 8\n"
		"\n"
		"buf: resq 1\n");
}

void readsource(FILE *f,char **sourcep){
	int sz=1024;
	char *source=malloc(sz);
	if(!source)outofmem();
	int i=0;
	while(true){
		if(i==sz-1){
			sz*=2;
			source=realloc(source,sz);
			if(!source)outofmem();
		}
		char c=fgetc(f);
		if(feof(f))break;
		if(!strchr("+-><[].,",c))continue;
		source[i++]=c;
	}
	source[i]='\0';
	*sourcep=source;
}

void optimise(char *source){
	int j=0;
#define COPY do {if(j<i)source[j]=source[i]; j++;} while(0)
	for(int i=0;source[i];i++){
		switch(source[i]){
			case '[':
				if(source[i+1]=='-'&&source[i+2]==']'){
					source[j++]='0';
					i+=2;
				} else COPY;
				break;

			default:
				COPY;
				break;
		}
	}
	source[j]='\0';
#undef COPY
}

int main(int argc,char **argv){
	fflush(stdout);
	parseargs(argc,argv);

	FILE *srcf=fopen(params.srcfname,"r");
	if(!srcf){
		fprintf(stderr,"Cannot open file '%s'\n",params.srcfname);
		return 1;
	}

	char *source;
	readsource(srcf,&source);
	assert(source);

	fclose(srcf);

	optimise(source);

	char asmfname[20];
	FILE *asmf=gettempfilew(asmfname);
	if(params.verbose)fprintf(stderr,"asm -> %s\n",asmfname);

	DBGTIME("writeprologue",writeprologue(asmf););
	DBGTIME("writeprogram",writeprogram(source,asmf););
	DBGTIME("writeepilogue",writeepilogue(asmf););

	fclose(asmf);

	/*char *b;
	asprintf(&b,"cat -n %s",asmfname);
	runcmd(b);
	free(b);*/
	int ret=compilechain(asmfname,params.dstfname);
	unlink(asmfname);
	if(params.verbose)fprintf(stderr,"unlink(%s)  (asmfname)\n",asmfname);
	return ret;
}