#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

#define MBR_SETTINGS_SIG_0 ((os_u16)0x4952)
#define MBR_SETTINGS_SIG_1 ((os_u16)0x4953)

int main(int argc, char** argv) {
    if (argc < 3) {
        fprintf(stderr, "Not enough arguments.\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.");

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

    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);

    os_bool found = OS_FALSE;
    os_s8* ptr = mbr;
    for(int i = 0; i < 512 - sizeof(MBR_BuildSettings); i++) {
        MBR_BuildSettings* settings = (MBR_BuildSettings*)ptr;
        if (settings->sig_0 == MBR_SETTINGS_SIG_0 && settings->sig_1 == MBR_SETTINGS_SIG_1) {
            settings->sector = kernel_offset / fat.header.base.bytes_per_sector;
            settings->size = kernel_size;
            found = OS_TRUE;
            break;
        }
        ptr++;
    }

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

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

    fclose(image);
    return 0;
}