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

#include "prologue.snippet.asm.h"
#include "epilogue.snippet.asm.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"
		"  -d             Copies generated assembly to <srcfile>.asm\n"
		"  -h             Show help\n"
		"  -H <HMfname>   Calculate heatmap while executing, writing data to <HMfname>\n"
		"  -L             Output last cell reached\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 lastcell,verbose,copyasm;
	const char *heatmapfname;
} Params;
Params params={30000,NULL,NULL,false,false,false,NULL};

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 'd':
						params.copyasm=true;
						break;

					case 'h':
						usage(argv[0]);
						exit(0);

					case 'H':
						if(i==argc-1)usage1(argv[0]);
						params.heatmapfname=argv[i+1];
						skipnext=true;
						break;

					case 'L':
						params.lastcell=true;
						break;

					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;

					default:
						usage1(argv[0]);
				}
			}
		} 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]);
}

typedef struct Range{
	int from,to; //[from,to); empty range indicates no mapping
} Range;

typedef struct Sourcemap{
	int len,origlen;
	Range *maps;
} Sourcemap;

Sourcemap* sourcemap_init(int len){
	Sourcemap *sm=malloc(sizeof(Sourcemap));
	if(!sm)outofmem();
	sm->len=len;
	sm->origlen=-1;
	sm->maps=calloc(len,sizeof(Range));
	if(!sm->maps)outofmem();
	return sm;
}

void sourcemap_destroy(Sourcemap *sm){
	if(sm->maps)free(sm->maps);
	free(sm);
}

int runcmd(const char *cmd){
	clock_t start=clock();
	pid_t pid=fork();
	if(pid<0){
		fprintf(stderr,"Cannot fork!\n");
		exit(1);
	}
	if(pid==0){
		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);
	fclose(f);

	char *buf;

	asprintf(&buf,"nasm -f macho64 -w+all -o %s %s",ofname,asmfname);
	if(!buf)outofmem();
	int ret=runcmd(buf);
	free(buf);
	if(ret)return ret;

	asprintf(&buf,"gcc %s -o %s",ofname,dstfname);
	if(!buf)outofmem();
	ret=runcmd(buf);
	free(buf);

	unlink(ofname);
	if(params.verbose)fprintf(stderr,"unlink(%s)  (ofile)\n",ofname);
	return ret;
}

char hexdigit(int n){
	if(n<0)return '?';
	else if(n<10)return n+'0';
	else if(n<16)return n-10+'a';
	else return '?';
}

void writesettings(FILE *asmf,Sourcemap *sm){
	fprintf(asmf,"%%define PARAMS_MEMSIZE %d\n",params.memsize);
	if(params.lastcell)fprintf(asmf,"%%define OUTPUT_LAST_CELL\n");
	if(params.heatmapfname){
		assert(sm);
		char buf[5*(strlen(params.heatmapfname)+1)];
		int i;
		for(i=0;i==0||params.heatmapfname[i-1];i++){
			buf[5*i+0]='0';
			buf[5*i+1]='x';
			buf[5*i+2]=hexdigit(params.heatmapfname[i]/16);
			buf[5*i+3]=hexdigit(params.heatmapfname[i]%16);
			buf[5*i+4]=',';
		}
		buf[5*(i-1)+4]='\0';
		fprintf(asmf,"%%define HEATMAPFNAME %s\n",buf);
		fprintf(asmf,"%%define SOURCE_LENGTH %d\n",sm->origlen);
	}
	fputc('\n',asmf);
}

void writeprologue(FILE *asmf){
	fwrite(prologue_snippet_asm,1,prologue_snippet_asm_len,asmf);
}

void writeprogram(char *source,FILE *asmf,Sourcemap *sm){
	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!='<'){
			if(params.lastcell){
				if(lastshiftc>0&&shiftc<0){
					int id=uniqid();
					fprintf(asmf,
						"%%ifdef OUTPUT_LAST_CELL\n"
						"\tcmp [maxp], rbx\n"
						"\tjge skip%d\n"
						"\tmov [maxp], rbx\n"
						"skip%d:\n"
						"%%endif\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;
		}

		if(sm){
			for(int j=sm->maps[i].from;j<sm->maps[i].to;j++){
				fprintf(asmf,"\tinc dword [r12+%d]\n",4*j);
			}
		}

		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){
	fwrite(epilogue_snippet_asm,1,epilogue_snippet_asm_len,asmf);
}

void readsource(FILE *f,char **sourcep,Sourcemap **smp){
	int sz=1024;
	Sourcemap *sm=NULL;
	if(smp)sm=sourcemap_init(sz);
	char *source=malloc(sz);
	if(!source)outofmem();
	int i=0;
	int sourcei=0;
	while(true){
		if(i==sz-1){
			sz*=2;
			source=realloc(source,sz);
			if(!source)outofmem();
			if(sm){
				sm->maps=realloc(sm->maps,sz*sizeof(Range));
				if(!sm->maps)outofmem();
				memset(sm->maps+sz/2,0,sz/2*sizeof(Range));
			}
		}
		char c=fgetc(f);
		if(feof(f))break;
		if(!strchr("+-><[].,",c)){
			sourcei++;
			continue;
		}
		if(sm){
			sm->maps[i].from=sourcei;
			sm->maps[i].to=sourcei+1;
		}
		source[i++]=c;
		sourcei++;
	}
	source[i]='\0';
	if(sm){
		sm->len=i;
		sm->origlen=sourcei;
	}
	*sourcep=source;
	if(sm)*smp=sm;
}

void optimise(char *source,const Sourcemap *srcsm,Sourcemap **dstsm){
	Sourcemap *sm=NULL;
	if(dstsm)sm=sourcemap_init(srcsm->len);
	int i,j=0;

#define COPY do { \
			if(j<i)source[j]=source[i]; \
			if(sm)sm->maps[j]=srcsm->maps[i]; \
			j++; \
		} while(0)

	for(i=0;source[i];i++){
		switch(source[i]){
			case '[':
				if(source[i+1]=='-'&&source[i+2]==']'){
					if(sm){
						sm->maps[j].from=srcsm->maps[i].from;
						sm->maps[j].to=srcsm->maps[i+2].to;
					}
					source[j++]='0';
					i+=2;
				} else COPY;
				break;

			default:
				COPY;
				break;
		}
	}
	if(sm){
		sm->len=j;
		sm->origlen=srcsm->origlen;
	}
	source[j]='\0';
#undef COPY
	if(sm)*dstsm=sm;
}

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;
	Sourcemap *sm1=NULL;
	readsource(srcf,&source,params.heatmapfname?&sm1:NULL);
	assert(source);

	fclose(srcf);

	Sourcemap *sm2=NULL;
	optimise(source,sm1,params.heatmapfname?&sm2:NULL);

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

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

	fclose(asmf);

	if(params.copyasm){
		char *b;
		asprintf(&b,"cp %s %s.asm",asmfname,params.srcfname);
		runcmd(b);
		free(b);
	}

	int ret=compilechain(asmfname,params.dstfname);
	unlink(asmfname);
	if(params.verbose)fprintf(stderr,"unlink(%s)  (asmfname)\n",asmfname);
	return ret;
}