#define MODULE "initrd"

#include "init/InitRD.h"
#include "bootboot.h"
#include "fs/VFS.h"
#include "io/Serial.h"
#include "log/Log.h"
#include "memory/MemoryManager.h"
#include "misc/utils.h"
#include "std/errno.h"
#include "std/stdlib.h"
#include "std/string.h"

extern BOOTBOOT bootboot;

static void* initrd_base;
static bool initrd_initialized = false;

static VFS::Node initrd_root;

bool InitRD::is_initialized()
{
    return initrd_initialized;
}

static inline uint64_t get_file_size_in_blocks(InitRD::File f)
{
    return f.size_in_blocks;
}

inline uint64_t InitRD::get_total_blocks()
{
    return bootboot.initrd_size / TAR_BLOCKSIZE;
}

inline InitRD::TarHeader* InitRD::get_block(uint64_t block_index)
{
    return (TarHeader*)((uintptr_t)initrd_base + block_index * TAR_BLOCKSIZE);
}

inline bool InitRD::is_valid_header(TarHeader* header)
{
    return strncmp(header->magic, TAR_MAGIC, 5) == 0;
}

uint64_t InitRD::get_file_physical_address(InitRD::File& file)
{
    return (uint64_t)file.addr - (uint64_t)initrd_base + (uint64_t)bootboot.initrd_ptr;
}

InitRD::File InitRD::get_file(TarHeader* header)
{
    File result;
    result.size = 0;
    memcpy(result.name, header->name, 100);
    int multiplier =
        1; // why they decided to store the size as an octal-encoded string instead of an integer is beyond me
    for (int i = 10; i >= 0; i--)
    {
        result.size += (multiplier * (header->size[i] - 48));
        multiplier *= 8;
    }
    result.addr = (void*)((uint64_t)header + TAR_BLOCKSIZE);
    result.size_in_blocks = Utilities::get_blocks_from_size(TAR_BLOCKSIZE, result.size);
    return result;
}

InitRD::File InitRD::open(const char* filename)
{
    uint64_t block = 0;
    uint64_t total_blocks = get_total_blocks();
    while (block < total_blocks)
    {
        TarHeader* hdr = (TarHeader*)get_block(block);
        if (hdr->typeflag == 53)
        {
            block++;
            continue;
        }
        if (!is_valid_header(hdr))
        {
            block++;
            continue;
        }
        auto f = get_file(hdr);
        if (strncmp(hdr->name, filename, strlen(filename)) == 0) { return f; }
        block += get_file_size_in_blocks(f) + 1;
    }
    File nullFile;
    nullFile.addr = 0;
    nullFile.size = 0;
    memcpy(nullFile.name, "NULL", 5);
    return nullFile;
}

void InitRD::for_each(void (*callback)(File& f))
{
    uint64_t block = 0;
    uint64_t total_blocks = get_total_blocks();
    while (block < total_blocks)
    {
        TarHeader* hdr = (TarHeader*)get_block(block);
        if (hdr->typeflag == 53)
        {
            block++;
            continue;
        }
        if (!is_valid_header(hdr))
        {
            block++;
            continue;
        }
        auto f = get_file(hdr);
        block += get_file_size_in_blocks(f) + 1;
        callback(f);
    }
}

#define INITRD_MAX_FILES_IN_DIR 16
#define INITRD_MAX_FILES 64

namespace InitRD
{
    struct Directory
    {
        char name[64];
        int entries = 0;
        VFS::Node* files[INITRD_MAX_FILES_IN_DIR];
    };
}

void initrd_for_each_dir(void (*callback)(InitRD::Directory& f))
{
    uint64_t block = 0;
    uint64_t total_blocks = InitRD::get_total_blocks();
    while (block < total_blocks)
    {
        InitRD::TarHeader* hdr = (InitRD::TarHeader*)InitRD::get_block(block);
        if (!InitRD::is_valid_header(hdr))
        {
            block++;
            continue;
        }
        if (hdr->typeflag == 53)
        {
            InitRD::Directory dir;
            strncpy(dir.name, hdr->name, sizeof(dir.name));
            callback(dir);
            block++;
            continue;
        }
        auto f = get_file(hdr);
        block += get_file_size_in_blocks(f) + 1;
    }
}

static InitRD::File files[INITRD_MAX_FILES];
static uint32_t total_files = 0;

static InitRD::Directory dirs[INITRD_MAX_FILES];
static uint32_t total_dirs = 0;

static VFS::Node nodes[INITRD_MAX_FILES + (INITRD_MAX_FILES - 1)]; // One of the dirs is the initrd_root
static uint32_t total_nodes = 0;

ssize_t initrd_read(VFS::Node* node, size_t offset, size_t length, char* buffer)
{
    if (!node) return -1;
    if (node->inode >= total_files) return -1;
    InitRD::File& file = files[node->inode];
    if (offset > file.size) return -1;
    if (offset + length > file.size) { length = file.size - offset; }
    memcpy(buffer, (void*)((uintptr_t)file.addr + offset), length);
    return length;
}

VFS::Node* initrd_scan_dir(VFS::Node* node, const char* filename)
{
    if (!node) return 0;
    if (node->inode >= total_dirs) return 0;
    InitRD::Directory dir = dirs[node->inode];
    for (int i = 0; i < dir.entries; i++)
    {
        if (strncmp(dir.files[i]->name, filename, sizeof(VFS::Node::name)) == 0) { return dir.files[i]; }
    }
    return 0;
}

VFS::Node* initrd_read_dir(VFS::Node* node, long offset)
{
    if (!node) return 0;
    if (node->inode >= total_dirs) return 0;
    InitRD::Directory dir = dirs[node->inode];
    if (offset >= dir.entries) return 0;
    return dir.files[offset];
}

int initrd_mkdir(VFS::Node* node, const char* name) // FIXME: Return proper error numbers.
{
    if (total_dirs >= 32)
    {
        kwarnln("mkdir() failed: too many directories");
        return -ENOSPC;
    }
    if (node->inode > total_dirs)
    {
        kwarnln("mkdir() failed: invalid node");
        return -EINVAL;
    }
    if (!(node->type & VFS_DIRECTORY))
    {
        kwarnln("mkdir() failed: not a directory");
        return -ENOTDIR;
    }
    InitRD::Directory& parent = dirs[node->inode];
    if (parent.entries == INITRD_MAX_FILES_IN_DIR)
    {
        kwarnln("mkdir() failed: parent is full");
        return -ENOSPC;
    }
    uint64_t inode = total_dirs;
    VFS::Node& new_node = nodes[total_nodes++];
    new_node.inode = inode;
    new_node.find_func = initrd_scan_dir;
    new_node.mkdir_func = initrd_mkdir;
    new_node.length = 0;
    new_node.type = VFS_DIRECTORY;
    strncpy(new_node.name, name, sizeof(new_node.name));
    InitRD::Directory dir;
    strncpy(dir.name, name, sizeof(dir.name));
    dir.entries = 0;
    dirs[total_dirs++] = dir; // FIXME: Right now this isn't of worry, but there is a possibility for a TOCTOU bug here.
                              // Should use a spinlock or something.
    node->length++;
    parent.files[parent.entries++] = &new_node;
    return 0;
}

static bool initrd_register_dir(InitRD::Directory& dir, uint64_t inode)
{
    const char* filename = dir.name;
    VFS::Node* current_node = &initrd_root;
    while (true)
    {
        while (*filename == '/') { filename++; }
        if (*filename == 0) { return false; }

        size_t path_section_size = 0;
        while (filename[path_section_size] && filename[path_section_size] != '/') { path_section_size++; }

        if (filename[path_section_size]) // We are in a '/'
        {
            char* buffer = (char*)kmalloc(path_section_size + 1);
            memcpy(buffer, filename, path_section_size);
            buffer[path_section_size] = 0;
            if (!current_node->find_func) { return false; }
            VFS::Node* child = current_node->find_func(current_node, buffer);
            if (!child) { return false; }
            current_node = child;
            kfree(buffer);
        }
        else
        {
            if (strncmp(filename, ".", path_section_size) != 0) // The current path section is not '.'
            {
                if (strncmp(filename, "..", path_section_size) == 0) { return false; }

                if (!current_node->find_func) { return false; }

                InitRD::Directory& parent = dirs[current_node->inode];
                if (parent.entries == INITRD_MAX_FILES_IN_DIR) { return false; }

                char* buffer = (char*)kmalloc(path_section_size + 1);
                memcpy(buffer, filename, path_section_size);
                buffer[path_section_size] = 0;

                VFS::Node& node = nodes[total_nodes++];
                node.inode = inode;
                node.find_func = initrd_scan_dir;
                node.length = 0;
                node.type = VFS_DIRECTORY;
                node.mkdir_func = initrd_mkdir;
                node.readdir_func = initrd_read_dir;
                node.length = 0;
                strncpy(node.name, buffer, sizeof(node.name));
                strncpy(dir.name, buffer, sizeof(dir.name));

                parent.files[parent.entries++] = &node;
                current_node->length++;

                kfree(buffer);
                return true;
            }
            else { return false; }
        }

        filename += path_section_size;
    }
}

static bool initrd_register_file(InitRD::File& f, uint64_t inode)
{
    const char* filename = f.name;
    VFS::Node* current_node = &initrd_root;
    while (true)
    {
        while (*filename == '/') { filename++; }
        if (*filename == 0) { return false; }

        size_t path_section_size = 0;
        while (filename[path_section_size] && filename[path_section_size] != '/') { path_section_size++; }

        if (filename[path_section_size]) // We are in a '/'
        {
            char* buffer = (char*)kmalloc(path_section_size + 1);
            memcpy(buffer, filename, path_section_size);
            buffer[path_section_size] = 0;
            if (!current_node->find_func) { return false; }
            VFS::Node* child = current_node->find_func(current_node, buffer);
            if (!child) { return false; }
            current_node = child;
            kfree(buffer);
        }
        else
        {
            if (strncmp(filename, ".", path_section_size) != 0) // The current path section is not '.'
            {
                if (strncmp(filename, "..", path_section_size) == 0) { return false; }

                if (!current_node->find_func) { return false; }

                InitRD::Directory& parent = dirs[current_node->inode];
                if (parent.entries == INITRD_MAX_FILES_IN_DIR) { return false; }

                char* buffer = (char*)kmalloc(path_section_size + 1);
                memcpy(buffer, filename, path_section_size);
                buffer[path_section_size] = 0;

                VFS::Node& node = nodes[total_nodes++];
                node.inode = inode;
                node.read_func = initrd_read;
                node.length = f.size;
                node.type = VFS_FILE;
                strncpy(node.name, buffer, sizeof(node.name));
                strncpy(f.name, buffer, sizeof(f.name));

                parent.files[parent.entries++] = &node;
                current_node->length++;
                kfree(buffer);
                return true;
            }
            else { return false; }
        }

        filename += path_section_size;
    }
}

static void initrd_scan()
{
    initrd_for_each_dir([](InitRD::Directory& dir) {
        if (total_dirs >= 32)
        {
            kwarnln("Failed to register directory %s: Too many directories in initrd", dir.name);
            return;
        }
        uint64_t inode = total_dirs;
        if (initrd_register_dir(dir, inode)) dirs[total_dirs++] = dir;
    });
    InitRD::for_each([](InitRD::File& f) {
        if (total_files >= 32)
        {
            kwarnln("Failed to register file %s: Too many files in initrd", f.name);
            return;
        }
        uint64_t inode = total_files;
        if (initrd_register_file(f, inode)) files[total_files++] = f;
    });
}

static void initrd_initialize_root()
{
    initrd_root.length = 0;
    initrd_root.inode = 0;
    initrd_root.type |= VFS_DIRECTORY;
    InitRD::Directory& root = dirs[0];
    total_dirs++;
    strncpy(initrd_root.name, "initrd", sizeof(initrd_root.name));
    strncpy(root.name, "initrd", sizeof(root.name));
    initrd_root.find_func = initrd_scan_dir;
    initrd_root.mkdir_func = initrd_mkdir;
    initrd_root.readdir_func = initrd_read_dir;
}

void InitRD::init()
{
    initrd_base = MemoryManager::get_unaligned_mappings(
        (void*)bootboot.initrd_ptr, Utilities::get_blocks_from_size(PAGE_SIZE, bootboot.initrd_size));
    kdbgln("physical base at %lx, size %lx, mapped to %p", bootboot.initrd_ptr, bootboot.initrd_size, initrd_base);
    kdbgln("total blocks: %ld", get_total_blocks());
    void* leak = kmalloc(4); // leak some memory so that kmalloc doesn't continously allocate and free pages
    initrd_initialize_root();
    initrd_scan();
    VFS::mount_root(&initrd_root);
    initrd_initialized = true;
    kfree(leak);
}