#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#define min(a, b) (((a) < (b)) ? (a) : (b))
// argv[1] must be floppy image
// argv[2] must be name of boot com
// if you expect to see memory safety, close your editor NOW!

#define AH 4
#define AX 0
#define CX 1
#define DH 6
#define BX 3
#define SP 4
#define SI 6
#define DI 7

char* push_imm8(char* addr, int i) {
	*(int8_t*)addr = i;
	return(addr+1);
}

char* push_imm16(char *addr, int i) {
	*(int16_t*)addr = i;
	return(addr+2);
}

#define push_mov8(a, b, c) push_imm8(push_imm8(a, 0xB0+b), c);
char* push_mov16(char* a, int b, int c) {
	if (c) {
		return(push_imm16(push_imm8(a, 0xB8+b), c));
	} else {
		// if value is 0 for 16-bit reg, xor it with itself
		// saves a byte
		return(push_imm16(a, 0xC031 | (b << 11) | (b << 8)));
	}
}

char* push_jc(char* addr, char* target) {
	int diff = target - addr;
	return(push_imm8(push_imm8(addr, 0x72), diff-2));
}

char* push_jmp(char* addr, char* target) {
	int diff = target - addr;
	if (abs(diff)<0x70) {
		return(push_imm8(push_imm8(addr, 0xEB), diff-2));
	} else {
		return(push_imm16(push_imm8(addr, 0xE9), diff-3));
	}
}

int main(int argc, char** argv) {
	// query mshowfat for first and last cluster number
	// we expect the file to be contiguous
	char output[64], cmd[256]= "mshowfat -i ";
	strcat(cmd, argv[1]);
	strcat(cmd, " '::");
	strcat(cmd, argv[2]);
	strcat(cmd, "'");
	FILE* fd = popen(cmd, "r");
	fgets(output, 64, fd);
	int r = pclose(fd);
	if (r > 0) {
		perror("Unable to query mtools");
		exit(1);
	}
	// parse the two numbers out of the output
	// output is like "::kernel.bs <2-4>"
	char* p = strchr(output, '<');
	p++;
	int min = strtol(p,&p,10);
	int max = 0;
	if (*p=='-') {
		p++;
		max = strtol(p,&p,10);
		if (*p!='>') {
			perror("Unexpected result from mshowfat");
			exit(1);
		}
	} else {
		max = min;
	}
	if (*p!='>') {
		perror("Unexpected result from mshowfat");
		exit(1);
	} 
	fd = fopen(argv[1], "r+");
	if (fd < 0) {
		perror("Cant open disk image");
	}
	char bootsect[512];
	if (fread(&bootsect, 512, 1, fd) < 1) {
		perror("Failed to read bootsector");
		exit(1);
	}
	// unpack fields from bpb
	int ss = *((uint16_t*)(bootsect+0x0B));
	int cs = *((uint8_t*)(bootsect+0x0D));
	int rsc = *((uint16_t*)(bootsect+0x0E));
	int fn = *((uint8_t*)(bootsect+0x10));
	int rde = *((uint16_t*)(bootsect+0x11));
	int fs = *((uint16_t*)(bootsect+0x16));
	int spt = *((uint16_t*)(bootsect+0x18));
	int nos = *((uint16_t*)(bootsect+0x1A));

	int base = rsc + (fs * fn) + ((rde * 32) / ss);

	int min_s = base + (min - 2) * cs;
	int max_s = base + (max - 2) * cs + (cs - 1);
	int bx = 0x100;
	int es = 0x60;

	char* bs = (char*)&bootsect;
	char* ptr = bs+0x3E;
	char* putsfunc = ptr;

	// assemble puts function
	//ptr = push_imm16(ptr, 0xAC5E);
	//ptr = push_imm16(ptr, 0xC084);
	//ptr = push_imm16(ptr, 0x0874);
	//ptr = push_imm16(ptr, 0x0EB4);
	//ptr = push_imm16(ptr, 0xDB31);
	//ptr = push_imm16(ptr, 0x10CD);
	//ptr = push_imm16(ptr, 0xF3EB);
	//ptr = push_imm16(ptr, 0xE6FF);

	// assemble entry point to our main code
	push_jmp(bs, ptr);

	// assemble segment and stack setup
	ptr = push_mov16(ptr, AX, es);
	ptr = push_imm16(ptr, 0xD08E); // mov sp, ax
	ptr = push_mov16(ptr, SP, 0);

	// save two copies of our working segment on the stack
	ptr = push_imm8(ptr, 0x50 + AX); // push ax
	ptr = push_imm8(ptr, 0x50 + AX); // push ax

	// assemble relocation code
	ptr = push_mov16(ptr, AX, 0);
	ptr = push_imm16(ptr, 0xD88E); // mov ds, ax
	ptr = push_mov16(ptr, SI, 0x7C00);
	ptr = push_imm16(ptr, 0xC08E); // mov es, ax
	ptr = push_mov16(ptr, DI, (es-0x10) << 4);
	ptr = push_mov16(ptr, CX, 0x0100);
	ptr = push_imm16(ptr, 0xA5F3); // rep movsw (does the copy)

	ptr = push_imm8(ptr, 0xEA); // jump far
	ptr = push_imm16(ptr, ptr+4-bs+0x600);
	ptr = push_imm16(ptr, 0);

	// pop DS and ES from the stack
	ptr = push_imm8(ptr, 0x07); // pop es
	ptr = push_imm8(ptr, 0x1F); // pop ds

	while(min_s <= max_s) {
		int sector = (min_s % spt) + 1;
		int tmp = min_s / spt;
		int head = tmp % nos;
		int track = tmp / nos;
		int len = min(spt - sector, max_s - min_s)+1;

		// memorize jump address in case of error
		char* retry = ptr;
		// set up registers for drive reset
		ptr = push_imm8(push_imm8(ptr, 0xB0+AH), 0);
		ptr = push_imm16(ptr, 0x13CD);
		// set up registers for read call
		ptr = push_imm16(push_imm8(ptr, 0xB8+AX), 0x0200 | len);
		ptr = push_imm16(push_imm8(ptr, 0xB8+CX), sector | (track << 8));
		ptr = push_imm8(push_imm8(ptr, 0xB0+DH), head);
		ptr = push_imm16(push_imm8(ptr, 0xB8+BX), bx);
		// set up read call
		ptr = push_imm16(ptr, 0x13CD);
		// jump if carry to retry addr
		ptr = push_jc(ptr, retry);

		min_s+=len;
		bx=bx+(ss*len);
	}

	// emerg return addr
	ptr = push_mov16(ptr, AX, 0);
	ptr = push_imm8(ptr, 0x50 + AX);

	ptr = push_imm8(ptr, 0xEA); // jmp far
	ptr = push_imm16(ptr, 0x100);
	ptr = push_imm16(ptr, es);

	// emerg return int 0
	ptr = bs + 0x100;
	ptr = push_imm16(ptr, 0x18CD);

	// write back bootsector
	fseek(fd, 0, SEEK_SET);
	if (fwrite(&bootsect, 512, 1, fd) < 1) {
		perror("Failed to write back boot sector");
		exit(1);
	}
}