#include <drivers/cpio/cpio.h>

#include <lib/memory.h>
#include <lib/string.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 (!(_dir->flags & VFS_DIRECTORY)) return NULL;

    vfs_node_t* child = _dir->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_get_or_create(vfs_node_t* root_node, const char* path) {
    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) {
        char* next_tkn = strtok_r(NULL, "/", &sav);
        
        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 = next_tkn;
    }
    
    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* parent = cpio_get_or_create(cpio_root, filename);

            char* base_name = filename;
            for (int i = strlen(filename) - 1; i >= 0; i--) {
                if (filename[i] == '/') {
                    base_name = &filename[i + 1];
                    break;
                }
            }

            vfs_node_t* current = cpio_find(parent, base_name);
            if (!current) {
                current = kalloc(sizeof(vfs_node_t));
                memset((void*)current, 0, sizeof(vfs_node_t));
                strcpy(current->name, base_name);
                
                current->next_sibling = parent->first_child;
                parent->first_child = current;
            }

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