//
// Created by rick on 19-09-21.
//

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <myke/vfs/blockdev.h>
#include <myke/vfs/lfs/ext2.h>
#include <myke/vfs/vfs-driver.h>
#include <myke/libk/libk.h>
#include <sys/param.h>

// https://wiki.osdev.org/Ext2

#define EXT2_DRIVER_NAME "ext2"
#define MI(mount) ((ext2_mount_info_t *)(mount)->global_driver_data)
#define FDI(fd) ((ext2_fd_info_t *)(fd)->driver_data)
#define DIV_ROUND_UP(a, b) (((a) + ((b) - 1)) / (b))

typedef struct ext2_dgetents_state {
    uint32_t bp;
    void *current_data;
    void *current_data_pos;
    void *current_data_end;
    uint32_t cur;
    bool at_end_of_block;
} ext2_dgetents_state_t;

typedef struct ext2_fd_info {
    uint32_t inode_nr;
    ext2_inode_t *inode;
    size_t seek_position;
    ext2_dgetents_state_t *dgetents_state;
    size_t size;
    bool dgetents_done;
} ext2_fd_info_t;

#define EXT2_READ_OK 0
#define EXT2_READ_ENOMEM 1
#define EXT2_READ_IOERR 2
#define EXT2_READ_EOF 3

int ext2_read_blocks(vfs_mount_t *mount, uint32_t block, uint32_t num, void *target) {
    uint32_t lba = (block * MI(mount)->block_size) / BLOCK_DEV_LBA_SIZE;
    uint32_t num_lba = DIV_ROUND_UP(num * MI(mount)->block_size, BLOCK_DEV_LBA_SIZE);
    void *buffer = malloc(num_lba * BLOCK_DEV_LBA_SIZE);
    if (buffer == NULL) {
        return EXT2_READ_ENOMEM;
    }

    if (((block_device_t *) MI(mount)->block_dev)->access(
            (block_device_t *) MI(mount)->block_dev,
            BLOCK_DEV_DIRECTION_READ, lba, num_lba, buffer)
        != BLOCK_DEV_ACCESS_OK) {
        free(buffer);
        return EXT2_READ_IOERR; // todo
    }
    memcpy(target, buffer, MI(mount)->block_size);
    free(buffer);
    return EXT2_READ_OK;
}

ext2_bg_descriptor_t *ext2_get_bg_descriptor(vfs_mount_t *mount, uint32_t idx) {
    uint32_t bg_per_block = MI(mount)->block_size / sizeof(ext2_bg_descriptor_t);
    uint32_t block = idx / bg_per_block;
    if (MI(mount)->block_size != 1024) {
        k_panics("only 1024 supported for now");
    }
    block += 2;
    ext2_bg_descriptor_t *blocks = malloc(MI(mount)->block_size);
    if (blocks == NULL) {
        return NULL;
    }
    if (ext2_read_blocks(mount, block, 1, blocks) != EXT2_READ_OK) {
        free(blocks);
        return NULL;
    }
    ext2_bg_descriptor_t *res = malloc(sizeof(ext2_bg_descriptor_t));
    if (res == NULL) {
        free(blocks);
        return NULL;
    }
    uint32_t block_index = idx % bg_per_block;
    memcpy(res, &blocks[block_index], sizeof(ext2_bg_descriptor_t));
    free(blocks);
    return res;
}

ext2_bg_descriptor_t *ext2_get_bg_descriptor_for_inode(vfs_mount_t *mount, uint32_t inode) {
    uint32_t group_idx = (inode - 1) / MI(mount)->sb.no_inodes_per_group;
    if (group_idx > MI(mount)->no_block_groups) {
        return NULL; // todo report
    }
    ext2_bg_descriptor_t *result = ext2_get_bg_descriptor(mount, group_idx);
    if (result == NULL) {
        return NULL;
    }
    return result;
}

ext2_inode_t *ext2_get_inode(vfs_mount_t *mount, uint32_t inode) {
    if (inode > MI(mount)->sb.total_inodes) {
        return NULL; // todo report
    }
    ext2_bg_descriptor_t *descriptor = ext2_get_bg_descriptor_for_inode(mount, inode);
    if (descriptor == NULL) {
        return NULL;
    }

    uint8_t *inodes_in_group = malloc(MI(mount)->block_size);
    if (inodes_in_group == NULL) {
        free(descriptor);
        return NULL;
    }

    uint32_t inode_index = (inode - 1) % MI(mount)->sb.no_inodes_per_group;
    uint32_t block = (inode_index * MI(mount)->sb.size_inode) / MI(mount)->block_size;
    uint32_t inode_offset = (inode_index * MI(mount)->sb.size_inode) % MI(mount)->block_size;

    // todo check bitmap?
    if (ext2_read_blocks(mount, descriptor->start_block_inode + block, 1, inodes_in_group) != EXT2_READ_OK) {
        free(descriptor);
        return NULL;
    }
    free(descriptor);

    ext2_inode_t *node = malloc(sizeof(ext2_inode_t));
    memcpy(node, &inodes_in_group[inode_offset], sizeof(ext2_inode_t));

    free(inodes_in_group);
    return node;
}

int ext2_get_block_from_inode(vfs_mount_t *mount, vfs_fd_t *fd, void *target, uint32_t block_idx) {
    if (block_idx < 12) {
        uint32_t bp = FDI(fd)->inode->dbp[block_idx];
        if (bp == 0) {
            return EXT2_READ_EOF;
        }
        return ext2_read_blocks(mount, bp, 1, target);
    } else {
        k_not_implemented();
    }
}

#define EXT2_DGETENT_CONT 1
#define EXT2_DGETENT_FULL 2
#define EXT2_DGETENT_END 3
#define EXT2_DGETENT_ERR 4

int ext2_dgetent_next(vfs_mount_t *mount, vfs_fd_t *fd, ext2_dgetents_state_t *state, void *target, size_t offset,
                      size_t max_size) {
    if (state->bp == UINT32_MAX) {
        // new
        state->current_data = malloc(MI(mount)->block_size);
        state->current_data_pos = state->current_data;
        state->current_data_end = state->current_data + MI(mount)->block_size;
        state->bp = 0;
        state->at_end_of_block = false;
        int result = ext2_get_block_from_inode(mount, fd, state->current_data, state->bp);
        if (result != EXT2_READ_OK) {
            return result == EXT2_READ_EOF ? EXT2_DGETENT_END : EXT2_DGETENT_ERR;
        }
    }
    if (state->at_end_of_block == true) {
        state->bp += 1;
        int result = ext2_get_block_from_inode(mount, fd, state->current_data, state->bp);
        state->current_data_pos = state->current_data;
        if (result != EXT2_READ_OK) {
            return result == EXT2_READ_EOF ? EXT2_DGETENT_END : EXT2_DGETENT_ERR;
        }
    }
    ext2_dir_entry_t *ent = state->current_data_pos;
    if (VFS_DIRENT_SIZE(ent->name_size) > max_size) {
        return EXT2_DGETENT_FULL;
    }
    vfs_mk_dirent_record(&target[offset], ent->inode, offset, 0, &ent->name, ent->name_size);
    state->current_data_pos += ent->ent_size;
    if (state->current_data_pos >= state->current_data_end) {
        state->at_end_of_block = true;
    }
    return EXT2_DGETENT_CONT;
}

int ext2_vfs_mount(vfs_mount_t *mount) {
    ext2_mount_info_t *mount_info = block_dev_mount(mount->device, EXT2_DRIVER_NAME);
    if (mount_info == NULL) {
        return VFS_MOUNT_ERROR;
    }
    mount->global_driver_data = mount_info;
    return VFS_MOUNT_OK;
}

ext2_fd_info_t *ext2_get_inode_info(vfs_mount_t *mount, uint32_t inode) {
    ext2_inode_t *node = ext2_get_inode(mount, inode);
    if (node == NULL) {
        return NULL;
    }
    ext2_fd_info_t *driver_data = malloc(sizeof(ext2_fd_info_t));
    if (driver_data == NULL) {
        return NULL;
    }
    size_t size = node->size_l;
    if (MI(mount)->sb.features_optional & EXT2_FEATURE_RO_64BIT_SIZE && node->type == EXT2_INODE_TYPE_FILE) {
#if EXT2_FEAT_RO_64BIT_SIZE
        size += node->size_h << 32;
#else
        if (node->size_h != 0) {
            printf("WARN: file to large");
            free(driver_data);
            return NULL;
        }
#endif
    }
    driver_data->inode_nr = inode;
    driver_data->inode = node;
    driver_data->size = size;
    driver_data->seek_position = 0;
    driver_data->dgetents_state = NULL;
    driver_data->dgetents_done = false;
    return driver_data;
}

vfs_dirent_t *ext2_get_dirent_for_filename(vfs_mount_t *mount, vfs_fd_t *dir, const char *path) {
    size_t name_length = strlen(path);
    vfs_dirent_t *ent = malloc(VFS_DIRENT_MAX_SIZE);
    if (ent == NULL) {
        return NULL;
    }
    ext2_dgetents_state_t *state = malloc(sizeof(ext2_dgetents_state_t));
    if (state == NULL) {
        return NULL;
    }
    memset(state, 0, sizeof(ext2_dgetents_state_t));
    state->bp = UINT32_MAX;
    bool found = false;
    // todo hash algorithm
    while (true) {
        const int result = ext2_dgetent_next(mount, dir, state, ent, 0, VFS_DIRENT_MAX_SIZE);
        if (result == EXT2_DGETENT_FULL || result == EXT2_DGETENT_ERR) {
            break; // whut
        }
        if (result == EXT2_DGETENT_END) {
            // not found
            break; // not found
        }
        if (name_length != strlen(ent->name)) {
            continue;
        }
        if (strncmp(path, ent->name, name_length) != 0) {
            // not same name
            continue;
        }
        found = true;
        break;
    }
    free(state);
    if (!found) {
        free(ent);
        return NULL;
    }
    return ent;
}

int ext2_vfs_open(vfs_mount_t *mount, vfs_fd_t *dir, const char *path, int mode, vfs_fd_t *out) {
    if (dir == NULL && path == NULL) {
        ext2_fd_info_t *driver_data = ext2_get_inode_info(mount, EXT2_ROOT_INODE);
        if (driver_data == NULL) {
            return VFS_OPEN_ERROR;
        }
        out->driver_data = driver_data;
        return VFS_OPEN_OK;
    }
    vfs_dirent_t *dirent = ext2_get_dirent_for_filename(mount, dir, path);
    int open_result = VFS_OPEN_ERROR;
    if (dirent != NULL) {
        ext2_fd_info_t *driver_data = ext2_get_inode_info(mount, dirent->inode);
        if (driver_data != NULL) {
            out->driver_data = driver_data;
            open_result = VFS_OPEN_OK;
        }
    }
    free(dirent);
    return open_result;
}

int ext2_vfs_close(vfs_mount_t *mount, vfs_fd_t *fd) {
    if (FDI(fd)->dgetents_state != NULL) {
        free(FDI(fd)->dgetents_state);
    }
    free(fd);
    return VFS_CLOSE_OK;
}

int ext2_vfs_fstat(vfs_mount_t *mount, vfs_fd_t *fd, stat_t *target, int flags) {
    target->st_dev = 0;  // todo
    target->st_ino = FDI(fd)->inode_nr;
    target->st_mode = FDI(fd)->inode->type;
    target->st_nlink = FDI(fd)->inode->no_hard_links;
    target->st_uid = FDI(fd)->inode->uid;
    target->st_gid = FDI(fd)->inode->gid;
    target->st_size = FDI(fd)->size;
    target->st_blksize = MI(mount)->block_size;
    target->st_blocks = FDI(fd)->size / 512; // todo correct?
    target->st_atim.tv_sec = FDI(fd)->inode->last_access;
    target->st_mtim.tv_sec = FDI(fd)->inode->last_mod;
    target->st_ctim.tv_sec = FDI(fd)->inode->created;

    if (FDI(fd)->inode->type == EXT2_INODE_TYPE_BLOCK_DEV || FDI(fd)->inode->type == EXT2_INODE_TYPE_CHAR_DEV) {
        target->st_rdev = FDI(fd)->inode->dbp[0];
    }
    return 0;
}

int ext2_vfs_fread(vfs_mount_t *mount, vfs_fd_t *fd, void *target, size_t size) {
    size_t target_pos = 0;

    void *buffer = malloc(MI(mount)->block_size);
    if (buffer == NULL) {
        return VFS_READ_ERROR;
    }
    while (true) {
        uint32_t block = FDI(fd)->seek_position / MI(mount)->block_size;
        uint32_t bytes_read_in_block = FDI(fd)->seek_position % MI(mount)->block_size;
        uint32_t bytes_left_in_block = MI(mount)->block_size - bytes_read_in_block;
        size_t toread = MIN(MIN(bytes_left_in_block, size - target_pos), FDI(fd)->size - target_pos);

        if (ext2_get_block_from_inode(mount, fd, buffer, block) != EXT2_READ_OK) {
            free(buffer);
            return VFS_READ_ERROR;
        }

        memcpy(&target[target_pos], &buffer[block], toread);
        target_pos += toread;
        FDI(fd)->seek_position += toread;
        if (target_pos == size || FDI(fd)->seek_position >= FDI(fd)->size) {
            break;
        }
    }
    free(buffer);
    return target_pos;
}

int ext2_vfs_dgetent(vfs_mount_t *mount, vfs_fd_t *fd, void *target, size_t size) {
    if (FDI(fd)->dgetents_done == true) {
        return 0;
    }
    size_t offset = 0;
    size_t prev_offset = 0;
    if (FDI(fd)->dgetents_state == NULL) {
        FDI(fd)->dgetents_state = malloc(sizeof(ext2_dgetents_state_t));
        FDI(fd)->dgetents_state->bp = UINT32_MAX;
        FDI(fd)->dgetents_state->cur = 0;
    }

    while (true) {
        int result = ext2_dgetent_next(mount, fd, FDI(fd)->dgetents_state, target, offset, size - offset);
        if (result != EXT2_DGETENT_CONT) {
            if (result == EXT2_DGETENT_END) {
                if (FDI(fd)->dgetents_state->current_data != NULL) {
                    free(FDI(fd)->dgetents_state->current_data);
                    FDI(fd)->dgetents_state->current_data = NULL;
                }
                free(FDI(fd)->dgetents_state);
                FDI(fd)->dgetents_state = NULL;
                FDI(fd)->dgetents_done = true;
            }
            if (result == EXT2_DGETENT_FULL || result == EXT2_DGETENT_END) {
                ((vfs_dirent_t *) &target[prev_offset])->offset_next = size;
            }
            break;
        }
        size_t reclen = ((vfs_dirent_t *) &target[offset])->reclen;
        prev_offset = offset;
        offset += reclen;
    }
    return offset;
}

uint32_t ext2_get_no_bg(ext2_sb_t *sb) {
    uint32_t no_bg_by_blocks = DIV_ROUND_UP(sb->total_blocks, sb->no_blocks_per_group);
    uint32_t no_bg_by_inodes = DIV_ROUND_UP(sb->total_inodes, sb->no_inodes_per_group);
    if (no_bg_by_blocks != no_bg_by_inodes) {
        k_panics("No no match");
    }
    return no_bg_by_inodes;
}

uint8_t ext2_check_device(block_device_t *device) {
    ext2_sb_t *sb = malloc(EXT2_SB_SIZE);
    if (device->access(device,
                       BLOCK_DEV_DIRECTION_READ,
                       EXT2_SB_ADDR / BLOCK_DEV_LBA_SIZE,
                       EXT2_SB_SIZE / BLOCK_DEV_LBA_SIZE,
                       sb) != BLOCK_DEV_ACCESS_OK) {
        free(sb);
        return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
    }
    if (sb->ext2_signature != EXT2_SIGNATURE) {
        printf("Missing signature\n");
        return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
    }
    printf("Ext2 features:\n\tRequired: %lx\n\tOptional: %lx\n\tRo: %lx\n", sb->features_required,
           sb->features_optional, sb->features_ro);
    if (sb->features_required ^ EXT2_FEAT_REQ_SUPPORTED) {
        printf("Filesystem uses features not supported by implementation, %lx\n", sb->features_required);
        free(sb);
        return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
    }
    printf("ext2 valid\n");
    ext2_mount_info_t *mount_info = malloc(sizeof(ext2_mount_info_t));
    memcpy(&mount_info->sb, sb, sizeof(ext2_sb_t));
    mount_info->no_block_groups = ext2_get_no_bg(sb);
    mount_info->block_size = 1024 << sb->block_size;
    mount_info->fragment_size = 1024 << sb->fragment_size;
    mount_info->block_dev = device;
    device->driver_info = mount_info;
    free(sb);
    return BLOCK_DEV_DRIVER_CHECK_OK;
}

BLOCK_DEV_DRIVER(300) = {
        .name = EXT2_DRIVER_NAME,
        .check_device = ext2_check_device,
        .free_device = NULL,
};

VFS_DRIVER(300) = {
        .name = EXT2_DRIVER_NAME,
        .vfs_mount = ext2_vfs_mount,
        .open = ext2_vfs_open,
        .close = ext2_vfs_close,
        .fstat = ext2_vfs_fstat,
        .fread = ext2_vfs_fread,
        .dgetent = ext2_vfs_dgetent,
};