#include <drivers/cpio/cpio.h>

#include <lib/memory.h>

#define MAX_INITRD_FILES 64
static vfs_node_t cpio_nodes[MAX_INITRD_FILES];
static uint32_t cpio_node_count = 0;

static uint32_t parse_hex_ascii(const char *str, int len) {
    uint32_t result = 0;
    for (int i = 0; i < len; i++) {
        result <<= 4;
        if (str[i] >= '0' && str[i] <= '9')      result |= (str[i] - '0');
        else if (str[i] >= 'A' && str[i] <= 'F') result |= (str[i] - 'A' + 10);
        else if (str[i] >= 'a' && str[i] <= 'f') result |= (str[i] - 'a' + 10);
    }
    return result;
}

static vfs_node_t* cpio_find(vfs_node_t* _dir, const char* name) {
    if (!(_dir->flags & VFS_DIRECTORY)) return NULL;

    for (uint32_t i = 0; i < cpio_node_count; i++) {
        if (strcmp(cpio_nodes[i].name, name) == 0) {
            return &cpio_nodes[i];
        }
    }
    
    return NULL; // File not found
}

static uint32_t cpio_read(struct vfs_node* node, uint32_t offset, uint32_t size, uint8_t* buffer) {
    if (offset >= node->size) return 0;
    if (offset + size > node->size) {
        size = node->size - offset;
    }
    memcpy((uint8_t*)(node->data) + offset, buffer, size);
    return size;
}

vfs_node_t* cpio_read_fs(uint32_t _vaddr) {
    cpio_header_t* header = (cpio_header_t*)_vaddr;

    vfs_node_t* cpio_root = kalloc(sizeof(vfs_node_t));
    memset((uint8_t*)cpio_root, 0, sizeof(vfs_node_t));

    strcpy(cpio_root->name, "cpio_root");

    cpio_root->flags |= VFS_DIRECTORY;
    cpio_root->finddir = cpio_find;

    uint32_t cnt = 0;
    while (1) {
        if (strncmp(header->c_magic, "070701", 6) != 0 && 
            strncmp(header->c_magic, "070702", 6) != 0) {
            break;
        }

        uint32_t namesize = parse_hex_ascii(header->c_namesize, 8);
        uint32_t filesize = parse_hex_ascii(header->c_filesize, 8);
        uint32_t mode     = parse_hex_ascii(header->c_mode, 8);

        char* filename = (char*)header + sizeof(cpio_header_t);

        if (strcmp(filename, "TRAILER!!!") == 0) {
            break;
        }

        uint32_t file_data_offset = sizeof(cpio_header_t) + namesize;
        file_data_offset = (file_data_offset + 3) & ~3; // align to 4 bytes

        uint8_t* file_data = (uint8_t*)header + file_data_offset;

        if (strcmp(".", filename) != 0 && cnt < MAX_INITRD_FILES) {
            vfs_node_t* current = &cpio_nodes[cpio_node_count++];
            memset((uint8_t*)current, 0, sizeof(vfs_node_t));

            strcpy(current->name, filename);
            current->size = filesize;

            current->flags = VFS_FILE;
            current->read = cpio_read;
            current->data = file_data;
        }

        uint32_t next_header_offset = file_data_offset + filesize;
        next_header_offset = (next_header_offset + 3) & ~3;
        header = (cpio_header_t*)((uint8_t*)header + next_header_offset);
    }

    return cpio_root;
}