#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include "fatfs.h"
#include "disk.h"

#define CPM_RECORD_SIZE 128
#define CPM_RECORDS_PER_BLOCK 128
#define CPM_BLOCK_SIZE CPM_RECORD_SIZE*CPM_RECORDS_PER_BLOCK

struct DirtyBlock {
	uint32_t block;
	struct DirtyBlock* next;
	uint8_t size;
	uint8_t* data;
};

typedef struct {
	FATFS *fs;
	uint8_t command;
	uint8_t ready;
	uint32_t counter;
	uint32_t lba;
	uint8_t* buffer;
	uint8_t* directory;
	unsigned int size;
	struct DirtyBlock* dirty_block;
} Storage;

Storage storage;

uint8_t disk_read_byte() {
	uint8_t value = 0x00;
	if (storage.ready && storage.command == 0x20) {
		if (storage.counter < storage.size) {
			value = storage.buffer[storage.counter];
		}

		storage.counter++;

		if (storage.counter >= CPM_RECORD_SIZE) {
			storage.ready = 0;
		}
	}

	return value;
}

void disk_write_byte(uint8_t value) {
	// @todo What if for whatever reason we write less then a record to disk?
	if (storage.ready && storage.command == 0x30) {
		storage.buffer[storage.counter] = value;
		storage.counter++;

		if (storage.counter >= CPM_RECORD_SIZE) {
			char filename[128] = {0};
			uint32_t offset = 0;
			if (storage.lba == 0) {
				snprintf(filename, 128, "0:loader.bin");
				offset = storage.lba - 0;
			} else if (storage.lba >= 1 && storage.lba < 45) {
				snprintf(filename, 128, "0:cpm22.bin");
				offset = storage.lba - 1;
			} else if (storage.lba >= 45 && storage.lba < 51) {
				snprintf(filename, 128, "0:bios.bin");
				offset = storage.lba - 45;
			} else if (storage.lba >= 256 && storage.lba < 384) {
				offset = storage.lba - 256;

				for (uint8_t index = 0; index < 4; ++index) {
					char name[8];
					for (int j = 0; j < 8; ++j) {
						char c = storage.buffer[index*32 + 1 + j];
						if (c == ' ') {
							c = 0x00;
						}
						name[j] = c;
					}

					char ext[3];
					for (int j = 0; j < 3; ++j) {
						char c = storage.buffer[index*32 + 9 + j];
						if (c == ' ') {
							c = 0x00;
						}
						ext[j] = c;
					}

					char temp_filename[128];
					snprintf(temp_filename, 128, "0:A/%.8s.%.3s", name, ext);

					uint8_t user = storage.buffer[index*32 + 0];
					uint8_t user_old = storage.directory[offset*128 + index*32 + 0];

					FIL file;
					if (user != user_old) {
						// @todo For some reason we keep getting a file entry that changes
						// from 0 -> 0xe5 with no name
						if (user == 0xe5) {
							FRESULT fr = f_unlink(temp_filename);
							if (fr) {
								printf("File error: %i\n\r", fr);
								return;
							}
						} else if (user_old == 0xe5) {
							FRESULT fr = f_open(&file, temp_filename, FA_READ | FA_CREATE_ALWAYS);
							if (fr) {
								printf("File error: %i\n\r", fr);
								return;
							}
							f_sync(&file);
							f_close(&file);
						}
					}

					for (uint8_t i = 0; i < 8; ++i) {
						uint16_t allocation = storage.buffer[index*32 + 16 + i*2];
						allocation += storage.buffer[index*32 + 17 + i*2] << 8;

						if (allocation == 0) {
							break;
						}

						struct DirtyBlock* prev = NULL;
						struct DirtyBlock* current = storage.dirty_block;
						while (current) {
							if (allocation == current->block) {
								FRESULT fr = f_open(&file, temp_filename, FA_WRITE);
								if (fr) {
									printf("File error: %i\n\r", fr);
									return;
								}

								offset = i*CPM_RECORDS_PER_BLOCK;
								fr = f_lseek(&file, offset * CPM_RECORD_SIZE);
								if (fr) {
									printf("File error: %i\n\r", fr);
									return;
								}

								fr = f_write(&file, current->data, current->size * CPM_RECORD_SIZE, &storage.size);
								if (fr) {
									printf("File error: %i\n\r", fr);
									return;
								}

								fr = f_sync(&file);
								if (fr) {
									printf("File error: %i\n\r", fr);
									return;
								}

								fr = f_close(&file);
								if (fr) {
									printf("File error: %i\n\r", fr);
									return;
								}

								if (prev) {
									prev->next = current->next;
								} else {
									storage.dirty_block = NULL;
								}

								free(current->data);
								free(current);

								break;
							}

							prev = current;
							current = current->next;
						}
					}
				}

				memcpy(storage.directory + offset*CPM_RECORD_SIZE, storage.buffer, CPM_RECORD_SIZE);
			} else if (storage.lba >= 384) {
				uint32_t block = (storage.lba - 256) / CPM_RECORDS_PER_BLOCK;

				uint8_t found = 0;
				for (int entry = 0; entry < 127; ++entry) {
					uint8_t user = storage.directory[entry*32 + 0];
					if (user == 0xe5) {
						continue;
					}

					for (int i = 0; i < 8; ++i) {
						uint16_t allocation = storage.directory[entry*32 + 16 + 2*i];
						allocation += storage.directory[entry*32 + 17 + 2*i] << 8;

						if (allocation == block) {
							found = 1;
							offset = (storage.lba - 256 - block*CPM_RECORDS_PER_BLOCK) + i*CPM_RECORDS_PER_BLOCK;

							char name[8];
							for (int j = 0; j < 8; ++j) {
								char c = storage.directory[entry*32 + 1 + j];
								if (c == ' ') {
									c = 0x00;
								}
								name[j] = c;
							}

							char ext[3];
							for (int j = 0; j < 3; ++j) {
								char c = storage.directory[entry*32 + 9 + j];
								if (c == ' ') {
									c = 0x00;
								}
								ext[j] = c;
							}

							snprintf(filename, 128, "0:A/%.8s.%.3s", name, ext);

							break;
						}
					}

					if (found) {
						break;
					}
				}

				if (!found) {
					// Check if we already have a dirty block for this block
					struct DirtyBlock* dirty_block = storage.dirty_block;
					while (dirty_block) {
						if (dirty_block->block == block) {
							break;
						}
					}

					// @todo It appears that it updates the file entry after every block written
					// So we might not need a linked list, altough it probably is more robust this way
					if (!dirty_block) {
						dirty_block = malloc(sizeof(struct DirtyBlock));
						dirty_block->block = block;
						dirty_block->next = storage.dirty_block;
						storage.dirty_block = dirty_block;
						dirty_block->size = 0;
						dirty_block->data = malloc(CPM_BLOCK_SIZE);
					}

					offset = (storage.lba - 256 - block*CPM_RECORDS_PER_BLOCK);
					memcpy(dirty_block->data + offset * CPM_RECORD_SIZE, storage.buffer, CPM_RECORD_SIZE);

					if (offset+1 > dirty_block->size) {
						dirty_block->size = offset+1;
					}
				}
			}

			if (filename[0]) {
				FIL file;
				// @0todo Cache the currently opened file so we are not opening it every 128 bytes
				FRESULT fr = f_open(&file, filename, FA_WRITE);
				if (fr) {
					printf("File error: %i\n\r", fr);
					return;
				}

				fr = f_lseek(&file, offset * CPM_RECORD_SIZE);
				if (fr) {
					printf("File error: %i\n\r", fr);
					return;
				}

				fr = f_write(&file, storage.buffer, storage.counter, &storage.size);
				if (fr) {
					printf("File error: %i\n\r", fr);
					return;
				}

				fr = f_sync(&file);
				if (fr) {
					printf("File error: %i\n\r", fr);
					return;
				}

				fr = f_close(&file);
				if (fr) {
					printf("File error: %i\n\r", fr);
					return;
				}
			}

			storage.ready = 0;
		}
	}
}

void disk_receive_command(uint8_t command) {
	storage.command = command;
	storage.counter = 0;
	storage.size = 0;

	if (storage.command == 0x20) {
		char filename[128] = {0};
		uint32_t offset = 0;
		if (storage.lba == 0) {
			snprintf(filename, 128, "0:loader.bin");
			offset = storage.lba - 0;
		} else if (storage.lba >= 1 && storage.lba < 45) {
			snprintf(filename, 128, "0:cpm22.bin");
			offset = storage.lba - 1;
		} else if (storage.lba >= 45 && storage.lba < 51) {
			snprintf(filename, 128, "0:bios.bin");
			offset = storage.lba - 45;
		} else if (storage.lba >= 256 && storage.lba < 384) {
			offset = storage.lba - 256;
			storage.size = 128;
			for (int i = 0; i < CPM_RECORD_SIZE; ++i) {
				storage.buffer[i] = storage.directory[offset*CPM_RECORD_SIZE + i];
			}
		} else if (storage.lba >= 384) {
			uint32_t block = (storage.lba - 256) / CPM_RECORDS_PER_BLOCK;

			uint8_t found = 0;
			for (int entry = 0; entry < 127; ++entry) {
				uint8_t user = storage.directory[entry*32 + 0];
				if (user == 0xe5) {
					continue;
				}

				for (int i = 0; i < 8; ++i) {
					uint16_t allocation = storage.directory[entry*32 + 16 + 2*i];
					allocation += storage.directory[entry*32 + 17 + 2*i] << 8;

					if (allocation == 0) {
						break;
					}

					if (allocation == block) {
						found = 1;
						offset = (storage.lba - 256 - block*CPM_RECORDS_PER_BLOCK) + i*CPM_RECORDS_PER_BLOCK;

						char name[8];
						for (int j = 0; j < 8; ++j) {
							char c = storage.directory[entry*32 + 1 + j];
							if (c == ' ') {
								c = 0x00;
							}
							name[j] = c;
						}

						char ext[3];
						for (int j = 0; j < 3; ++j) {
							char c = storage.directory[entry*32 + 9 + j];
							if (c == ' ') {
								c = 0x00;
							}
							ext[j] = c;
						}

						snprintf(filename, 128, "0:A/%.8s.%.3s", name, ext);

						break;
					}
				}

				if (found) {
					break;
				}
			}
		}

		// @todo Maybe add a mechanism to read dirty blocks?

		if (filename[0]) {
			FIL file;
			// @todo Cache the currently opened file so we are not opening it every 128 bytes
			FRESULT fr = f_open(&file, filename, FA_READ);
			if (fr) {
				printf("File error: %i\n\r", fr);
				return;
			}

			fr = f_lseek(&file, offset * CPM_RECORD_SIZE);
			if (fr) {
				printf("File error: %i\n\r", fr);
				return;
			}

			fr = f_read(&file, storage.buffer, CPM_RECORD_SIZE, &storage.size);
			if (fr) {
				printf("File error: %i\n\r", fr);
				return;
			}

			fr = f_close(&file);
			if (fr) {
				printf("File error: %i\n\r", fr);
				return;
			}
		}

		storage.ready = 1;
	} else if (storage.command == 0x30) {
		// We don't need to prepare enything
		storage.ready = 1;
	}
}

uint8_t disk_is_ready() {
	return storage.ready;
}

void disk_get_entries() {
	// Clear all entries
	for (int i = 0; i < 0x4000; ++i) {
		if (i % 32 == 0) {
			storage.directory[i] = 0xe5;
		} else {
			storage.directory[i] = 0x00;
		}
	}

	DIR dir;
	FILINFO fno;
	FRESULT fr = f_opendir(&dir, "0:A");
	if (fr == FR_OK) {
		uint16_t end = 1;
		uint8_t index = 0;
		for (;;) {
			fr = f_readdir(&dir, &fno);
			if (fr != FR_OK || fno.fname[0] == 0) {
				break;
			}

			if (!(fno.fattrib & AM_DIR)) {
				printf("File: %s\n\r", fno.fname);
				storage.directory[index*32 + 0] = 0x00;

				// Calculate the number of records
				uint32_t size = (fno.fsize + CPM_RECORD_SIZE - 1) / CPM_RECORD_SIZE;

				storage.directory[index*32 + 16] = end & 0xFF;
				storage.directory[index*32 + 17] = (end >> 8) & 0xFF;
				end++;

				int i = 1;
				while (size >= 0x80) {
					size -= 0x80;
					storage.directory[index*32 + 16 + i*2] = end & 0xFF;
					storage.directory[index*32 + 17 + i*2] = (end >> 8) & 0xFF;
					end++;
					i++;
				}

				storage.directory[index*32 + 12] = (i-1) & 0xFF;
				storage.directory[index*32 + 13] = 0x00;
				storage.directory[index*32 + 14] = ((i-1) >> 8) & 0xFF;

				// Store the size in the number of records
				storage.directory[index*32 + 15] = size;
				/* free(entries[index].filename); */

				uint8_t name_len = 9;
				for (uint8_t i = 0; i < 8; ++i) {
					char c = ' ';

					if (name_len == 9 && fno.fname[i] != '.') {
						c = fno.fname[i];
					} else if (name_len == 9){
						name_len = i+1;
					}

					storage.directory[index*32 + 1+i] = c;
				}

				uint8_t done = 0;
				for (uint8_t i = 0; i < 3; ++i) {
					char c = ' ';
					
					if (!done && fno.fname[name_len+i] != 0) {
						c = fno.fname[name_len+i];
					} else {
						done = 1;
					}

					storage.directory[index*32 + 9+i] = c;
				}

				index++;
			}
		}
	} else {
		printf("File error: %i\n\r", fr);
		return;
	}


	fr = f_closedir(&dir);
	if (fr) {
		printf("File error: %i\n\r", fr);
		return;
	}

}

void disk_init() {
	storage.command = 0;
	storage.ready = 0;
	storage.counter = 0;
	storage.lba = 0;
	storage.size = 0;
	storage.dirty_block = NULL;

	storage.buffer = (uint8_t*)malloc(CPM_RECORD_SIZE);
	storage.directory = (uint8_t*)malloc(0x4000);

	storage.fs = (FATFS*)malloc(sizeof(FATFS));

	FRESULT fr = f_mount(storage.fs, "0:", 0);
	if (fr) {
		printf("File error: %i\n\r", fr);
	}

	disk_get_entries();
}

void disk_free() {
	free(storage.buffer);
	free(storage.directory);

	f_mount(0, "0:", 0);
	free(storage.fs);

	while (storage.dirty_block) {
		struct DirtyBlock* next = storage.dirty_block->next;
		free(storage.dirty_block->data);
		free(storage.dirty_block);
		storage.dirty_block = next;
	}
}

void disk_set_lba_part(uint8_t value, uint8_t part) {
	uint32_t temp = storage.lba & ~(0xFF << 8*part);
	storage.lba = temp + (value << 8*part);
}