#include <drivers/cpio/cpio.h>

#include <lib/memory.h>
#include <lib/string.h>
#include <lib/print.h>

#define MAX_INITRD_FILES 64

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 (strcmp(name, "tvga") == 0) {
        kprintf("looking for tvga\n");
    }
    if (!(_dir->flags & VFS_DIRECTORY)) return NULL;
    vfs_node_t* root = _dir;
    if (_dir->mnt) {
        root = _dir->mnt;
    }

    vfs_node_t* child = root->first_child; 
    while (child != NULL) {
        if (strcmp(child->name, name) == 0) {
            return child;
        }
        child = child->next_sibling;
    }
    
    return NULL;
}

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

static vfs_node_t* cpio_find_or_create(vfs_node_t* root_node, const char* path) {
    if (strcmp(path, "fonts") == 0) {
        kprintf("found fonts folder\n");
    }
    vfs_node_t* current = root_node;
    
    char path_cpy[256];
    strcpy(path_cpy, path);
    
    char* sav = NULL;
    char* tkn = strtok_r(path_cpy, "/", &sav);
    
    while (tkn != NULL) {
        vfs_node_t* next_node = cpio_find(current, tkn);
        if (!next_node) {
            next_node = kalloc(sizeof(vfs_node_t));
            memset((void*)next_node, 0, sizeof(vfs_node_t));
            
            strcpy(next_node->name, tkn);
            next_node->flags = VFS_DIRECTORY;
            next_node->find = cpio_find;
            
            next_node->next_sibling = current->first_child;
            current->first_child = next_node;
        }
        
        current = next_node;
        tkn = strtok_r(NULL, "/", &sav);
    }
    
    return current;
}

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->find = cpio_find;

    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 + 3) & ~3;
        uint8_t* file_data = (uint8_t*)header + file_data_offset;

        if (strcmp(".", filename) != 0) {
            vfs_node_t* current = cpio_find_or_create(cpio_root, filename);
            if ((mode & 0xF000) == 0x4000) {
                current->flags = VFS_DIRECTORY;
                current->find = cpio_find;
            } else {
                current->flags = VFS_FILE;
                current->size = filesize;
                current->data = file_data;
                current->read = cpio_read;
            }
        }

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