/**
 * An attempt to write a tool for ext4 that will allow to change inode count of a filesystem.
 *
 * In theory it shouldn't be that hard:
 * 1) If shrinking:
 * 1.1) mark the end of each inode table as reserved (to make sure the inode allocator won't allocate it)
 * 1.2) if there were some inodes, move them away:
 * 1.2.1) allocate a new place for each inode, copy it there
 * 1.2.2) remember the old->new inode number mapping
 * 1.3) mark some blocks that belonged to inode table as free
 * 2) If growing:
 * 2.1) check all extra blocks that will be occupied by the growing inode tables,
 *      mark them as occupied if there are free ()
 * 2.2) move data away from 'extra blocks' that were occupied
 * 3) Change all inode numbers in directory entries so that
 *    new_num = (old_num/old_inodes_per_group)*new_inodes_per_group + (old_num%inodes_per_group)
 *    also translate old->new inode numbers remembered at (1.2.2)
 * 4) Change superblock: s_inodes_count, s_free_inodes_count, s_inodes_per_group
 * 5) Change block group descriptors: bg_inode_table_(lo,hi), bg_free_inodes_count_(lo,hi),
 *    bg_inode_bitmap_csum_(lo,hi), bg_itable_unused_(lo,hi)
 * 6) If flex_bg is enabled, move parts of old inode tables so they are consecutive again
 *
 * This is a destructive process involving metadata change so it would be good if
 * we could first write a file containing all necessary changes, then backup all
 * changed blocks, and then write new blocks to the disk. (maybe just use undo_io_manager?)
 */

#include <stdio.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/time.h>
#include <errno.h>

#include <ext2fs/ext2_fs.h>
#include <ext2fs/ext2fs.h>

#define _(a) (a)

// "local data" for the inode reallocation process
typedef struct
{
    ext2_filsys fs;
    int fs_fd;
    char *device_name, *io_options;
    __u32 new_inode_count, new_inodes_per_group, new_inode_blocks_per_group;
    // (old->new) inode number map
    ext2_ino_t *inode_map;
    __u32 inode_map_size, inode_map_alloc;
} realloc_data;

// Utility functions for (old -> new) inode number map
void realloc_add_inode_map(realloc_data *rd, ext2_ino_t old, ext2_ino_t new)
{
    if (!old || !new)
    {
        return;
    }
    if (2*rd->inode_map_size >= rd->inode_map_alloc)
    {
        rd->inode_map_alloc += 1024;
        rd->inode_map = realloc(rd->inode_map, sizeof(ext2_ino_t) * rd->inode_map_alloc);
    }
    rd->inode_map[rd->inode_map_size*2] = old;
    rd->inode_map[rd->inode_map_size*2+1] = new;
    rd->inode_map_size++;
}

int realloc_compare_inode_map_callback(const void *a, const void *b)
{
    return *((ext2_ino_t*)a) - *((ext2_ino_t*)b);
}

void realloc_sort_inode_map(realloc_data *rd)
{
    if (!rd->inode_map)
    {
        return;
    }
    qsort(rd->inode_map, rd->inode_map_size, sizeof(ext2_ino_t)*2, realloc_compare_inode_map_callback);
}

ext2_ino_t realloc_search_inode_map(realloc_data *rd, ext2_ino_t old)
{
    __u32 start = 0, end = rd->inode_map_size, cur;
    ext2_ino_t cur_ino;
    if (!rd->inode_map)
    {
        return 0;
    }
    while (end-start > 1)
    {
        cur = (start+end)>>1;
        cur_ino = rd->inode_map[cur<<1];
        if (cur_ino < old)
        {
            start = cur+1;
        }
        else if (cur_ino > old)
        {
            end = cur;
        }
        else
        {
            return rd->inode_map[(cur<<1)+1];
        }
    }
    return 0;
}

/**
 * Move inodes from the end of each block group inode table
 * so the tables can be shrinked
 */
int shrink_move_inodes(realloc_data *rd)
{
    int retval = 0, inode_size = EXT2_INODE_SIZE(rd->fs->super);
    __u32 group, i;
    __u32 new_group;
    ext2_ino_t ino, new_ino;
    struct ext2_inode *inode = NULL;
    ext2fs_read_inode_bitmap(rd->fs);
    if (retval)
    {
        return retval;
    }
    retval = ext2fs_get_mem(inode_size, &inode);
    if (retval)
    {
        return retval;
    }
    for (group = 0; group < rd->fs->group_desc_count; group++)
    {
        for (i = rd->new_inodes_per_group; i < EXT2_INODES_PER_GROUP(rd->fs->super); i++)
        {
            ino = 1 + group*EXT2_INODES_PER_GROUP(rd->fs->super) + i;
            if (ext2fs_test_inode_bitmap2(rd->fs->inode_map, ino))
            {
                // Inode is occupied and should be moved
                new_group = group;
                do
                {
                    retval = ext2fs_find_first_zero_inode_bitmap2(rd->fs->inode_map,
                        1 + new_group*EXT2_INODES_PER_GROUP(rd->fs->super),
                        1 + new_group*EXT2_INODES_PER_GROUP(rd->fs->super)+rd->new_inodes_per_group, &new_ino);
                    if (!retval)
                    {
                        break;
                    }
                    new_group = (new_group+1) % rd->fs->group_desc_count;
                } while (new_group != group);
                if (retval)
                {
                    // No space to move this inode
                    goto out;
                }
                // Copy inode to the new place
                retval = ext2fs_read_inode_full(rd->fs, ino, inode, inode_size);
                if (retval)
                {
                    goto out;
                }
                retval = ext2fs_write_inode_full(rd->fs, new_ino, inode, inode_size);
                if (retval)
                {
                    goto out;
                }
                retval = ext2fs_mark_inode_bitmap2(rd->fs->inode_map, new_ino);
                if (retval)
                {
                    goto out;
                }
                ext2fs_mark_ib_dirty(rd->fs);
                // Remember mapping
                realloc_add_inode_map(rd, ino, new_ino);
            }
        }
    }
out:
    if (inode)
    {
        ext2fs_free_mem(&inode);
    }
    return retval;
}

/**
 * Move data blocks from after the end of each block group inode table
 * so the tables can be grown
 */
int extend_move_blocks(realloc_data *rd)
{
    
    return ENOSYS;
}

static int change_inode_numbers_callback(ext2_ino_t dir, int entry,
    struct ext2_dir_entry *dirent, int offset,
    int blocksize, char *buf, void *priv_data)
{
    realloc_data *rd = priv_data;
    ext2_ino_t new_ino = realloc_search_inode_map(rd, dirent->inode);
    if (!new_ino)
    {
        new_ino = dirent->inode;
    }
    new_ino = 1 + (new_ino-1)/EXT2_INODES_PER_GROUP(rd->fs->super)*rd->new_inodes_per_group +
        (new_ino-1)%EXT2_INODES_PER_GROUP(rd->fs->super);
    if (new_ino != dirent->inode)
    {
        dirent->inode = new_ino;
        return DIRENT_CHANGED;
    }
    return 0;
}

/**
 * Change inode numbers in all directory entries
 */
int change_inode_numbers(realloc_data *rd)
{
    ext2_ino_t ino;
    realloc_sort_inode_map(rd);
    for (ino = 1; ino <= rd->fs->super->s_inodes_count; ino++)
    {
        ext2fs_dir_iterate2(rd->fs, ino, 0, 0, change_inode_numbers_callback, rd);
    }
    return 0;
}

/**
 * 1) Move inode tables so they are consecutive again if flex_bg is enabled
 * 2) Mark/unmark extra inode table blocks
 * 3) Adjust superblock and block group descriptors
 */
int change_super_and_bgd(realloc_data *rd)
{
    blk64_t it_start, blk;
    dgrp_t grp, n_flex, n_grp;
    __u32 unus;
    int i_per_g_diff = 0;
    int flexbg_size = 0, i, retval = 0;
    void *buf = NULL;
    ext2fs_flush(rd->fs);
    ext2fs_read_block_bitmap(rd->fs);
    if (rd->new_inode_blocks_per_group != rd->fs->inode_blocks_per_group)
    {
        // Move inode tables if flex_bg is active
        if (EXT2_HAS_INCOMPAT_FEATURE(rd->fs->super, EXT4_FEATURE_INCOMPAT_FLEX_BG)
            && rd->fs->super->s_log_groups_per_flex)
        {
            flexbg_size = 1 << rd->fs->super->s_log_groups_per_flex;
            retval = ext2fs_get_mem(EXT2_BLOCK_SIZE(rd->fs->super) * rd->new_inode_blocks_per_group * flexbg_size, &buf);
            if (retval)
            {
                goto out;
            }
            i_per_g_diff = rd->new_inodes_per_group - EXT2_INODES_PER_GROUP(rd->fs->super);
            for (n_flex = 0; n_flex < rd->fs->group_desc_count/flexbg_size; n_flex++)
            {
                n_grp = flexbg_size;
                if (n_flex*flexbg_size+n_grp > rd->fs->group_desc_count)
                {
                    n_grp = rd->fs->group_desc_count-n_flex*flexbg_size;
                }
                // Read inode tables
                for (grp = n_flex*flexbg_size, i = 0; i < n_grp; grp++, i++)
                {
                    if (!ext2fs_bg_flags_test(rd->fs, grp, EXT2_BG_INODE_UNINIT))
                    {
                        blk = ext2fs_inode_table_loc(rd->fs, grp);
                        retval = io_channel_read_blk64(rd->fs->io, blk, rd->new_inode_blocks_per_group,
                            buf + i*rd->new_inode_blocks_per_group*EXT2_BLOCK_SIZE(rd->fs->super));
                        if (retval)
                        {
                            goto out;
                        }
                    }
                }
                // Write inode table to the new place
                it_start = ext2fs_inode_table_loc(rd->fs, n_flex*flexbg_size);
                blk = rd->new_inode_blocks_per_group * n_grp;
                retval = io_channel_write_blk64(rd->fs->io, it_start, blk, buf);
                if (retval)
                {
                    // The FS is badly corrupted now :-(
                    printf("Error moving inode tables for %u groups, starting from %u\n", n_grp, n_flex*flexbg_size);
                    goto out;
                }
                // Mark/unmark extra inode table blocks
                if (rd->new_inode_blocks_per_group < rd->fs->inode_blocks_per_group)
                {
                    ext2fs_unmark_block_bitmap_range2(rd->fs->block_map, it_start+blk,
                        (rd->fs->inode_blocks_per_group-rd->new_inode_blocks_per_group)*n_grp);
                }
                else
                {
                    ext2fs_mark_block_bitmap_range2(rd->fs->block_map, it_start+blk,
                        (rd->new_inode_blocks_per_group-rd->fs->inode_blocks_per_group)*n_grp);
                }
                if (!ext2fs_bg_flags_test(rd->fs, n_flex*flexbg_size, EXT2_BG_INODE_UNINIT))
                {
                    ext2fs_bg_free_blocks_count_set(rd->fs, n_flex*flexbg_size,
                        ext2fs_bg_free_blocks_count(rd->fs, n_flex*flexbg_size) -
                        (rd->new_inode_blocks_per_group - rd->fs->inode_blocks_per_group)*flexbg_size);
                }
                // Change inode table locations and free inode counts
                for (grp = n_flex*flexbg_size, i = 0; i < n_grp; grp++, i++)
                {
                    blk = it_start + rd->new_inode_blocks_per_group*i;
                    ext2fs_inode_table_loc_set(rd->fs, grp, blk);
                    ext2fs_bg_free_inodes_count_set(rd->fs, grp,
                        ext2fs_bg_free_inodes_count(rd->fs, grp) + i_per_g_diff);
                    unus = ext2fs_bg_itable_unused(rd->fs, grp);
                    unus = unus < i_per_g_diff ? 0 : unus - i_per_g_diff;
                    ext2fs_bg_itable_unused_set(rd->fs, grp, unus);
                    ext2fs_group_desc_csum_set(rd->fs, grp);
                }
            }
        }
        else
        {
            // Mark/unmark extra inode table blocks (without flex_bg)
            for (grp = 0; grp < rd->fs->group_desc_count; grp++)
            {
                it_start = ext2fs_inode_table_loc(rd->fs, grp);
                if (rd->new_inode_blocks_per_group < rd->fs->inode_blocks_per_group)
                {
                    ext2fs_unmark_block_bitmap_range2(rd->fs->block_map, it_start,
                        rd->fs->inode_blocks_per_group-rd->new_inode_blocks_per_group);
                }
                else
                {
                    ext2fs_mark_block_bitmap_range2(rd->fs->block_map, it_start,
                        rd->new_inode_blocks_per_group-rd->fs->inode_blocks_per_group);
                }
                ext2fs_bg_free_blocks_count_set(rd->fs, grp, ext2fs_bg_free_blocks_count(rd->fs, grp) -
                    rd->new_inode_blocks_per_group + rd->fs->inode_blocks_per_group);
                ext2fs_bg_free_inodes_count_set(rd->fs, grp, ext2fs_bg_free_inodes_count(rd->fs, grp) +
                    rd->new_inodes_per_group - EXT2_INODES_PER_GROUP(rd->fs->super));
                unus = ext2fs_bg_itable_unused(rd->fs, grp);
                unus = unus < i_per_g_diff ? 0 : unus - i_per_g_diff;
                ext2fs_bg_itable_unused_set(rd->fs, grp, unus);
                ext2fs_group_desc_csum_set(rd->fs, grp);
            }
        }
    }
    ext2fs_mark_bb_dirty(rd->fs);
    retval = rd->fs->write_bitmaps(rd->fs);
    if (retval)
    {
        goto out;
    }
    rd->fs->write_bitmaps = NULL;
    ext2fs_free_blocks_count_add(rd->fs->super, rd->fs->group_desc_count * (rd->fs->inode_blocks_per_group - rd->new_inode_blocks_per_group));
    rd->fs->super->s_free_inodes_count += rd->fs->group_desc_count * (rd->new_inodes_per_group - EXT2_INODES_PER_GROUP(rd->fs->super));
    rd->fs->super->s_inodes_per_group = rd->new_inodes_per_group;
    rd->fs->super->s_inodes_count = rd->fs->group_desc_count * rd->new_inodes_per_group;
    ext2fs_mark_super_dirty(rd->fs);
out:
    if (buf)
    {
        ext2fs_free_mem(&buf);
    }
    return retval;
}

/**
 * Main function: change inode number of a filesystem!
 */
int do_realloc(realloc_data *rd)
{
    int retval;
    rd->new_inodes_per_group = rd->new_inode_count / rd->fs->group_desc_count;
    rd->new_inodes_per_group &= ~7;
    if (rd->new_inodes_per_group < 16)
    {
        printf("Too small number of inodes requested (%u), min inodes per group = 16\n", rd->new_inodes_per_group);
        return ENOENT;
    }
    rd->new_inode_blocks_per_group =
        (rd->new_inodes_per_group * EXT2_INODE_SIZE(rd->fs->super) +
        EXT2_BLOCK_SIZE(rd->fs->super) - 1) / EXT2_BLOCK_SIZE(rd->fs->super);
    rd->new_inode_count = rd->new_inodes_per_group * rd->fs->group_desc_count;
    if (rd->new_inodes_per_group < EXT2_INODES_PER_GROUP(rd->fs->super))
    {
        if (rd->new_inode_count < rd->fs->super->s_inodes_count - rd->fs->super->s_free_inodes_count)
        {
            printf("Too small number of inodes requested, existing inodes (%u) won't fit\n",
                rd->fs->super->s_inodes_count - rd->fs->super->s_free_inodes_count);
            return ENOENT;
        }
        printf("Phase 1: Moving inodes out of the way\n");
        retval = shrink_move_inodes(rd);
        if (retval)
        {
            return retval;
        }
    }
    else if (rd->new_inodes_per_group > EXT2_INODES_PER_GROUP(rd->fs->super))
    {
        blk64_t required_blocks = (rd->new_inode_blocks_per_group - rd->fs->inode_blocks_per_group) * rd->fs->group_desc_count;
        if (required_blocks > ext2fs_free_blocks_count(rd->fs->super))
        {
            printf("Requested number of inodes is too big, it requires %llu free blocks, "
                "and there are only %llu free blocks available\n",
                required_blocks, ext2fs_free_blocks_count(rd->fs->super));
            return ENOENT;
        }
        printf("Phase 1: Moving data blocks out of the way\n");
        retval = extend_move_blocks(rd);
        if (retval)
        {
            return retval;
        }
    }
    else
    {
        printf("The requested number of inodes is equal to current\n");
        return 0;
    }
    printf("Phase 2: Changing all inode numbers\n");
    retval = change_inode_numbers(rd);
    if (retval)
    {
        return retval;
    }
    printf("Phase 3: Adjusting superblock and block group descriptors\n");
    retval = change_super_and_bgd(rd);
    if (retval)
    {
        return retval;
    }
    return 0;
}

__u32 atou(char *s)
{
    __u32 x = 0;
    if (s[0] == '0')
    {
        if (s[1] == 'x' || s[1] == 'X')
        {
            sscanf(s+2, "%x", &x);
        }
        else
        {
            sscanf(s+1, "%o", &x);
        }
    }
    else
    {
        sscanf(s, "%u", &x);
    }
    return x;
}

const char *program_name = "realloc-inodes";

int main(int narg, char **args)
{
    realloc_data rd = { 0 };
    int optind, retval, io_flags = 0, force = 0;
    ext2fs_struct_stat st_buf;
    if (narg < 3)
    {
        printf("USAGE: ./realloc-inodes <device> <new_inode_count>\n");
        return 0;
    }
    optind = 1;
    rd.device_name = args[optind++];
    rd.new_inode_count = atou(args[optind++]);
    add_error_table(&et_ext2_error_table);
    // Open FS
    rd.fs_fd = ext2fs_open_file(rd.device_name, O_RDWR, 0);
    if (rd.fs_fd < 0)
    {
        com_err("open", errno, _("while opening %s"), rd.device_name);
        exit(1);
    }
    retval = ext2fs_fstat(rd.fs_fd, &st_buf);
    if (retval < 0)
    {
        com_err("open", errno, _("while getting stat information for %s"), rd.device_name);
        exit(1);
    }
    if (!S_ISREG(st_buf.st_mode))
    {
        close(rd.fs_fd);
        rd.fs_fd = -1;
    }
    rd.io_options = strchr(rd.device_name, '?');
    if (rd.io_options)
    {
        *rd.io_options++ = 0;
    }
    io_flags = EXT2_FLAG_64BITS | EXT2_FLAG_RW | EXT2_FLAG_EXCLUSIVE;
    retval = ext2fs_open2(rd.device_name, rd.io_options, io_flags, 0, 0, unix_io_manager, &rd.fs);
    if (retval)
    {
        com_err(program_name, retval, _("while trying to open %s"), rd.device_name);
        printf(_("Couldn't find valid filesystem superblock.\n"));
        goto close_fd;
    }
    if (!force && ((rd.fs->super->s_state & EXT2_ERROR_FS) || ((rd.fs->super->s_state & EXT2_VALID_FS) == 0)))
    {
        fprintf(stderr, _("Please run 'e2fsck -f %s' first.\n\n"), rd.device_name);
        goto close_fs;
    }
    // Call main realloc function
    retval = do_realloc(&rd);
    if (retval)
    {
        com_err(program_name, retval, _("while resizing inode count"));
        goto close_fs;
    }
close_fs:
    ext2fs_close(rd.fs);
close_fd:
    if (rd.fs_fd > 0)
    {
        close(rd.fs_fd);
    }
    return retval;
}