#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/stdlib.h"
#include <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);
    }
}

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

static VFS::Node nodes[32];

int32_t initrd_read(VFS::Node* node, uint32_t offset, uint32_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 = (uint32_t)file.size - offset; }
    memcpy(buffer, (void*)((uint64_t)file.addr + offset), length);
    return (int32_t)length;
}

VFS::Node* initrd_scan_root(VFS::Node*, const char* filename)
{
    for (uint32_t i = 0; i < total_files; i++)
    {
        if (strncmp(nodes[i].name, filename, sizeof(VFS::Node::name)) == 0) { return &nodes[i]; }
    }
    return 0;
}

void initrd_scan()
{
    InitRD::for_each([](InitRD::File& f) {
        if (total_files >= 32) return;
        kdbgln("registering file %s", f.name);
        files[total_files++] = f;
    });
    for (uint32_t i = 0; i < total_files; i++)
    {
        VFS::Node& node = nodes[i];
        node.inode = i;
        node.read_func = initrd_read;
        node.length = files[i].size;
        strncpy(node.name, files[i].name, sizeof(node.name));
    }
}

void InitRD::init()
{
    initrd_base =
        MemoryManager::get_unaligned_mappings((void*)bootboot.initrd_ptr, bootboot.initrd_size / PAGE_SIZE + 1);
    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());
    initrd_scan();
    initrd_root.length = 0;
    initrd_root.inode = (uint64_t)-1;
    strncpy(initrd_root.name, "initrd", 7);
    initrd_root.find_func = initrd_scan_root;
    VFS::mount_root(&initrd_root);
    VFS::mount_root(0);
    VFS::mount_root(&initrd_root);
    kdbgln("mounted initrd at VFS root, total files: %d", total_files);
    initrd_initialized = true;
}