#include "libfat.h"
#include "libhost.h"
#include "stdio.h"
#include "string.h"
#include "stdlib.h"

os_u32 fat_buffer[512];

#define STAGE1_NAME "OS3ST1  "
#define STAGE1_EXT "BIN"
#define STAGE1_FULL_NAME "OS3ST1  BIN"

OS_PACK_START
typedef struct _MBR_BuildSettings {
    os_u16 sig_0;
    os_u16 sector;
    os_u16 size;
    os_u16 sig_1;
} OS_PACK_MID MBR_BuildSettings;
OS_PACK_END

OS_PACK_START
typedef struct _MBR_ReadPacket {
	os_u16 sig_0;
	os_u8 packet_size; //Fixed 10h
	os_u8 reserved;
	os_u16 blocks;
	os_u16 buffer_offset;
	os_u16 buffer_segment;
	//os_u32 buffer;
	os_u64 lba;
	os_u16 sig_1;
} OS_PACK_MID MBR_ReadPacket;
OS_PACK_END

#define MBR_SETTINGS_SIG_0 ((os_u16)0x4952)
#define MBR_SETTINGS_SIG_1 ((os_u16)0x4953)
#define MBR_PACKET_SIG_0 ((os_u16) 0x4954)
#define MBR_PACKET_SIG_1 ((os_u16) 0x4955)

int main(int argc, char** argv) {
    if (argc < 3) {
        fprintf(stderr, "Not enough arguments. Usage: kernelcopy <image file> <kernel file>\n");
        return 1;
    }

    FILE* image = fopen(argv[1], "rb+");
    FILE* kernel = fopen(argv[2], "rb");

    if (image == NULL) {
        fprintf(stderr, "Failed to open image file\n");
        return 1;
    }

    LFFileHandle image_handle = libhost_create_handle(image);

    LFControlBlock fat;
    if(libfat_init(image_handle, fat_buffer, sizeof(fat_buffer), &fat) != OS_TRUE) {
        fprintf(stderr, "Failed to read FAT header\n");
        return 1;
    }

    printf("FAT OK.\n");

    LFDirectory root_dir;
    if(libfat_open_directory(&fat, 0, &root_dir) != OS_TRUE) {
        fprintf(stderr, "Failed to open root directory\n");
        return 1;
    }

    LFFile handle;
    os_bool handle_found = OS_FALSE;
    LFDirectoryEntry entry;
    while(libfat_read_directory(&root_dir, &entry) == OS_TRUE) {
        if(strcmp(entry.entry.file_name, STAGE1_NAME) == 0 && strcmp(entry.entry.extension, STAGE1_EXT) == 0) {
            printf("Found existing 'OS3ST1.BIN'.\n");
            if(libfat_truncate(&fat, &entry) != OS_TRUE) {
                fprintf(stderr, "Failed to truncate existing kernel.\n");
                return 1;
            }

            handle_found = OS_TRUE;
            break;
        }
    }

    if(handle_found == OS_FALSE) {
        if(libfat_create(&fat, 0, STAGE1_FULL_NAME, OS_FALSE, &entry) != OS_TRUE) {
            fprintf(stderr, "Failed to create kernel file\n");
            return 1;
        }
        printf("Created 'OS3ST1.BIN'.\n");
    }


    if(libfat_open(&fat, &entry, &handle) != OS_TRUE) {
        fprintf(stderr, "Failed to open kernel.\n");
        return 1;
    }


    char* buffer = malloc(8192);
    os_u32 kernel_size = 0;
    os_u32 read;
    do {
        read = fread(buffer, 1, 8192, kernel);
        if(read > 0) {
            if(libfat_write(&handle, buffer, read) != read) {
                free(buffer);
                fprintf(stderr, "Error during kernel copy\n");
                return  1;
            }
        }
        kernel_size += read;
    } while(read > 0);
    printf("Kernel copy OK.\n");

    free(buffer);
    fclose(kernel);

    os_u32 kernel_offset = libfat_get_true_entry_offset(&fat, &handle.file_entry);

    if (kernel_offset == 0) {
        fprintf(stderr, "Failed to get true file sector.\n");
        return 1;
    }

	printf("Bytes per sector: %d\n", fat.header.base.bytes_per_sector);
	printf("Kernel offset: %d, difference from bytes per sector: %d\n", kernel_offset, kernel_offset % fat.header.base.bytes_per_sector);
	fflush(stdout);
    if (kernel_offset % fat.header.base.bytes_per_sector != 0) {
        fprintf(stderr, "Error: STAGE1 not aligned to sector\n");
        return 1;
    }

    fseek(image, 0, SEEK_SET);
    char mbr[512];
    fread(mbr, 1, 512, image);

	printf("Kernel size: %d", kernel_size);

	os_bool found = OS_FALSE;
	os_u8* ptr = mbr;
	for (int i = 0; i < 512 - sizeof(MBR_ReadPacket); i++) {
		MBR_ReadPacket * packet = (MBR_ReadPacket*)ptr;
		if(packet->sig_0 == MBR_PACKET_SIG_0 && packet->sig_1 == MBR_PACKET_SIG_1) {
			packet->lba = kernel_offset / fat.header.base.bytes_per_sector;
			packet->blocks = (kernel_size / fat.header.base.bytes_per_sector) + ((kernel_size % fat.header.base.bytes_per_sector) == 0 ? 0 : 1);
			packet->buffer_offset = 0x0000;
			packet->buffer_segment = 0x1000;
			//packet->buffer = 0x10000;
			found = OS_TRUE;
			break;
		}
		ptr++;
	}

	if (found == OS_FALSE) {
		fprintf(stderr, "Error: could not find READ_PACKET markers\n");
		return 1;
	}

    fseek(image, 0, SEEK_SET);
    fwrite(mbr, 1, 512, image);

    fclose(image);
    return 0;
}