#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string.h>
#include <sys/types.h>
#include "control.h"
#include "io.h"
#include "bsp_driver_sd.h"
#include "profiling.h"

// The top 4 address bits determine which device is used. (16 pages, 256 devices)
// 0xFF means that the device does not exist
// WATCH OUT THE DEVICE ADDRESS NEED TO BE REVERSED (MSB IS ON THE RIGHT)
uint8_t memory_map_0[16] = {0b00010000, 0b10001000, 0b01001000, 0b11001000, 0b00101000, 0b10101000, 0b01101000, 0b11101000, 0b00011000, 0b10011000, 0b01011000, 0b11011000, 0b00111000, 0b10111000, 0b01111000, 0b11111000};
uint8_t memory_map_1[16] = {0b00001000, 0b10001000, 0b01001000, 0b11001000, 0b00101000, 0b10101000, 0b01101000, 0b11101000, 0b00011000, 0b10011000, 0b01011000, 0b11011000, 0b00111000, 0b10111000, 0b01111000, 0b11111000};

Control control;

extern UART_HandleTypeDef huart2;
extern uint8_t ack[1];
extern uint8_t nack[1];

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

uint8_t get_device(uint16_t address) {
	uint8_t page = address >> 12;
	if (control.memory_config == 0) {
		return memory_map_0[page];
	} else if (control.memory_config == 1) {
		return memory_map_1[page];
	}

	return 0xFF;
}

void control_program_eeprom(uint8_t* data, uint16_t length) {
	// Take control of the bus
	send_busrq(1);
	while (!has_busak()) {
		control_cycle();
	}

	enable_address_out(1);

	select_device(memory_map_0[0]);

	for (uint16_t i = 0; i < length; ++i) {
		write_address(i);
		write_data(data[i]);

		enable_data_out(1);
		send_memrq(1);
		send_wr(1);

		send_wr(0);
		enable_data_out(0);
		send_rd(1);

		for (;;) {
			uint8_t d = read_data();
			if (d == data[i]) {
				break;
			}
		}

		send_rd(0);
		send_memrq(0);

		uint8_t progress[] = {(i+1) & 0xFF, (i+1) >> 8};
		HAL_UART_Transmit(&huart2, progress, sizeof(progress), HAL_MAX_DELAY);
	}

	enable_address_out(0);

	// Release the bus again
	send_busrq(0);

	// Restart the z80
	control_reset();

	HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);
}

void handle_memrq() {
	uint16_t address = read_address();
	uint8_t device = get_device(address);
	select_device(device);
}

void handle_io_read() {
	uint8_t address = read_address() & 0xFF;

	switch (address) {
		// @todo This should be detected on startup
		// Stand in for graphics hardware
		/* case 0x03: */
		/* 	write_data(0x01); */
		/* 	break; */

		/* Stand in for the keyboard hardware */
		/* case 0x1E: */
		/* 	write_data(char_c); */
		/* 	char_r = 0; */
		/* 	break; */

		/* Stand in for the keyboard hardware */
		/* case 0x1F: */
		/* 	#<{(| write_data(0x01 * char_r); |)}># */
		/* 	write_data(0x00); */
		/* 	break; */

		// Read byte from disk
		case 0x08:
			if (control.storage.ready && control.storage.command == 0x20) {
				if (control.storage.counter < control.storage.size) {
					write_data(control.storage.buffer[control.storage.counter]);
				} else {
					write_data(0x00);
				}

				control.storage.counter++;

				if (control.storage.counter >= CPM_RECORD_SIZE) {
					control.storage.ready = 0;
				}
			} else {
				write_data(0x00);
			}
			break;

		// Check if disk is ready
		case 0x0f:
			write_data(0x08*control.storage.ready);
			break;

		default: {
			/* uint8_t value = read_data(); */
			/* #<{(| if (value == 0) { |)}># */
			/* 	printf("IO Read: %.2X @ %.2X\n\r", value, address); */
			/* #<{(| } |)}># */
			return;
		}
	}

	enable_data_out(1);
}

void handle_io_write() {
	uint8_t address = read_address() & 0xFF;
	uint8_t value = read_data();

	switch (address) {
		case 0x00:
			control.memory_config = 0;
			break;

		case 0x01:
			control.memory_config = 1;
			break;

		case 0x02:
			printf("%c", value);
			break;

		// Write byte to disk
		case 0x08:
			// @todo What if for whatever reason we write less then a record to disk?
			if (control.storage.ready && control.storage.command == 0x30) {
				control.storage.buffer[control.storage.counter] = value;
				control.storage.counter++;

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

						for (uint8_t index = 0; index < 4; ++index) {
							char name[8];
							for (int j = 0; j < 8; ++j) {
								char c = control.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 = control.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 = control.storage.buffer[index*32 + 0];
							uint8_t user_old = control.storage.directory[offset*128 + index*32 + 0];

							FIL file;
							if (user != user_old) {
								printf("User has changed for %s: %i -> %i\n\r", temp_filename, user_old, user);

							// @todo For some reason we keep getting a file entry that changes
							// from 0 -> 0xe5 with no name
								if (user == 0xe5) {
									printf("Index: %i\n\r", index);
									printf("Deleting file\n\r");
									FRESULT fr = f_unlink(temp_filename);
									if (fr) {
										printf("File error: %i\n\r", fr);
										return;
									}
								} else if (user_old == 0xe5) {
									printf("Creating file\n\r");

									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 = control.storage.buffer[index*32 + 16 + i*2];
								allocation += control.storage.buffer[index*32 + 17 + i*2] << 8;

								if (allocation == 0) {
									break;
								}

								struct DirtyBlock* prev = NULL;
								struct DirtyBlock* current = control.storage.dirty_block;
								while (current) {
									if (allocation == current->block) {
										printf("Found dirty block\n\r");

										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, &control.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 {
											control.storage.dirty_block = NULL;
										}

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

										break;
									}

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

						printf("Write to directory (offset: %li)\n\r", offset);
						memcpy(control.storage.directory + offset*CPM_RECORD_SIZE, control.storage.buffer, CPM_RECORD_SIZE);
					} else if (control.storage.lba >= 384) {
						uint32_t block = (control.storage.lba - 256) / CPM_RECORDS_PER_BLOCK;

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

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

								if (allocation == block) {
									found = 1;
									offset = (control.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 = control.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 = control.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) {
							printf("Write to unlinked block\n\r");

							// Check if we already have a dirty block for this block
							struct DirtyBlock* dirty_block = control.storage.dirty_block;
							while (dirty_block) {
								if (dirty_block->block == block) {
									printf("Found existing dirty block\n\r");
									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) {
								printf("Creating new dirty block\n\r");
								dirty_block = malloc(sizeof(struct DirtyBlock));
								dirty_block->block = block;
								dirty_block->next = control.storage.dirty_block;
								control.storage.dirty_block = dirty_block;
								dirty_block->size = 0;
								dirty_block->data = malloc(CPM_BLOCK_SIZE);
							}

							offset = (control.storage.lba - 256 - block*CPM_RECORDS_PER_BLOCK);
							memcpy(dirty_block->data + offset * CPM_RECORD_SIZE, control.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, control.storage.buffer, control.storage.counter, &control.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;
						}
					}

					control.storage.ready = 0;
				}
			}
			break;

		case 0x0b: {
			uint32_t temp = control.storage.lba & 0xFFFF00;
			control.storage.lba = temp + value;
			break;
		}

		case 0x0c: {
			uint32_t temp = control.storage.lba & 0xFF00FF;
			control.storage.lba = temp + (value << 8);
			break;
		}

		case 0x0d: {
			uint32_t temp = control.storage.lba & 0x00FFFF;
			control.storage.lba = temp + (value << 16);
			break;
		}

		// Receive disk command
		case 0x0f: {
			control.storage.command = value;
			control.storage.counter = 0;
			control.storage.size = 0;

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

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

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

							if (allocation == 0) {
								break;
							}

							if (allocation == block) {
								found = 1;
								offset = (control.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 = control.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 = control.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, control.storage.buffer, CPM_RECORD_SIZE, &control.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;
					}
				}

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

		default:
			printf("IO Write: %.2X @ %.2X\n\r", value, address);
			break;
	}
}

void handle_ioreq() {
		if (has_wr()) {
			handle_io_write();
		} else if (has_rd()) {
			handle_io_read();
		}
}

void control_cycle() {
	set_clock(1);

	// We need this not detect IO multiple times
	static uint8_t had_ioreq = 0;
	if (!has_ioreq()) {
		had_ioreq = 0;
	}

	// @todo We are forgetting to set this somewhere
	enable_data_out(0);

	if (has_memrq()) {
		handle_memrq();
	} else if (has_ioreq() && !has_m1()) {
		had_ioreq++;
		if (had_ioreq == 3) {
			handle_ioreq();
		}
	} else if (has_ioreq() && has_m1()) {
		printf("Interrupt ackknowledged\n\r");
	}

	set_clock(0);
}

void get_disk_entries(char* diskpath) {
	printf("Scanning: %s\n\r", diskpath);

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

	DIR dir;
	FILINFO fno;
	FRESULT fr = f_opendir(&dir, diskpath);
	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);
				control.storage.directory[index*32 + 0] = 0x00;

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

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

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

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

				// Store the size in the number of records
				control.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;
					}

					control.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;
					}

					control.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;
	}
}

// @todo Properly reset everything
void control_reset() {
	free(control.storage.buffer);
	free(control.storage.directory);

	f_close(control.storage.file);
	free(control.storage.file);

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

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

	Control temp = {0, {NULL, NULL, 0, 0, 0, 0, NULL, NULL, 0, NULL}};
	control = temp;

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

	control.storage.fs = (FATFS*)malloc(sizeof(FATFS));
	control.storage.file = (FIL*)malloc(sizeof(FIL));

	printf("Status: %i\n\r", disk_status(0));

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

	get_disk_entries("0:A");

	set_reset(1);
	for (int i = 0; i <= 10; ++i) {
		set_clock(i % 2);
	}
	set_reset(0);
}