#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define int8    int8_t
#define int16   int16_t
#define int32   int32_t
#define int64   int64_t
#define uint8   uint8_t
#define uint16  uint16_t
#define uint32  uint32_t
#define uint64  uint64_t

typedef int64 bfs_off_t;

struct block_run {
        int32           allocation_group;
        uint16          start;
        uint16          length;
};

typedef block_run inode_addr;

#define BFS_DISK_NAME_LENGTH    32

struct disk_super_block {
        char            name[BFS_DISK_NAME_LENGTH];
        int32           magic1;
        int32           fs_byte_order;
        uint32          block_size;
        uint32          block_shift;
        bfs_off_t       num_blocks;
        bfs_off_t       used_blocks;
        int32           inode_size;
        int32           magic2;
        int32           blocks_per_ag;
        int32           ag_shift;
        int32           num_ags;
        int32           flags;
        block_run       log_blocks;
        bfs_off_t       log_start;
        bfs_off_t       log_end;
        int32           magic3;
        inode_addr      root_dir;
        inode_addr      indices;
        int32           pad[8];
};

#define SUPER_BLOCK_MAGIC1                      'BFS1'          /* BFS1 */

const char *kHexDigits = "0123456789abcdef";


static void
read_from(int fd, off_t pos, void *buffer, size_t size)
{
        if (lseek(fd, pos, SEEK_SET) < 0
                || read(fd, buffer, size) != (int)size) {
                fprintf(stderr, "Error: Failed to read %lu bytes at offset %lld\n",
                        size, pos);
                exit(1);
        }
}


static void
write_to(int fd, off_t pos, const void *buffer, size_t size)
{
        if (lseek(fd, pos, SEEK_SET) < 0
                || write(fd, buffer, size) != (int)size) {
                fprintf(stderr, "Error: Failed to write %lu bytes at offset %lld\n",
                        size, pos);
                exit(1);
        }
}

static int
count_bits(int value)
{
        int count = 0;

        while (value) {
                if (value & 1)
                        count++;
                value >>= 1;
        }

        return count;
}

int
main(int argc, const char *const *argv)
{
        if (argc < 2 || argc > 3) {
                fprintf(stderr, "Usage: %s <image> [ <hex OR pattern> ]\n", argv[0]);
                exit(1);
        }

        const char *fileName = argv[1];
        const char *patternString = (argc >= 3 ? argv[2] : "0f");

        // skip leading "0x" in the pattern string

        if (strncmp(patternString, "0x", 2) == 0)
                patternString += 2;

        // translate the pattern

        uint8 *pattern = new uint8[strlen(patternString) / 2 + 1];
        int patternLen = 0;
        for (; *patternString; patternString++) {
                if (!isspace(*patternString)) {
                        char c = *patternString;
                        const char *digitPos = strchr(kHexDigits, tolower(c));
                        if (!digitPos) {
                                fprintf(stderr, "Error: Invalid pattern character: \"%c\"\n",
                                        c);
                                exit(1);
                        }
                        int digit = digitPos - kHexDigits;

                        if (patternLen & 1)
                                pattern[patternLen / 2] |= digit;
                        else
                                pattern[patternLen / 2] = digit << 4;
                        patternLen++;
                }
        }

        if (patternLen == 0 || (patternLen & 1)) {
                fprintf(stderr, "Error: Invalid pattern! Must have even number (>= 2) "
                        "of hex digits.\n");
                exit(1);
        }
        patternLen /= 2;        // is now length in bytes


        printf("bfs_fragmenter: using pattern: 0x");
        for (int i = 0; i < patternLen; i++)
                printf("%02x", pattern[i]);
        printf("\n");

        // open file

        int fd = open(fileName, O_RDWR);
        if (fd < 0) {
                fprintf(stderr, "Error: Failed to open \"%s\": %s\n", fileName,
                        strerror(errno));
                exit(1);
        }

        // read superblock

        disk_super_block superBlock;
        read_from(fd, 512, &superBlock, sizeof(superBlock));

        if (superBlock.magic1 != SUPER_BLOCK_MAGIC1) {
                fprintf(stderr, "Error: Bad superblock magic. This is probably not a "
                        "BFS image.\n");
                exit(1);
        }
        int blockSize = superBlock.block_size;

        // stat the image

        struct stat st;
        if (fstat(fd, &st) < 0) {
                fprintf(stderr, "Error: Failed to stat image: %s\n", strerror(errno));
                exit(1);
        }

        int64 blockCount = st.st_size / blockSize;
        if (blockCount != superBlock.num_blocks) {
                fprintf(stderr, "Error: Number of blocks in image and the number in "
                        "the superblock differ: %lld vs. %lld\n", blockCount,
                        superBlock.num_blocks);
                exit(1);
        }


        // iterate through the block bitmap blocks and or the bytes with 0x0f

        uint8 *block = new uint8[blockSize];
        int64 occupiedBlocks = 0;
        int64 bitmapByteCount = blockCount / 8;         // round down, we ignore the

                                                                                                // last partial byte

        int64 bitmapBlockCount = (bitmapByteCount + blockSize - 1) / blockSize;

        int bitmapByteIndex = 0;
        for (int i = 1; i < bitmapBlockCount + 1; i++) {
                off_t offset = i * blockSize;
                read_from(fd, offset, block, blockSize);

                int toProcess = blockSize;
                if (toProcess > bitmapByteCount - bitmapByteIndex)
                        toProcess = (bitmapByteCount - bitmapByteIndex);

                for (int k = 0; k < toProcess; k++) {
                        uint8 patternChar = pattern[bitmapByteIndex % patternLen];
                        occupiedBlocks += count_bits(~block[k] & patternChar);
                        block[k] |= patternChar;
                        bitmapByteIndex++;
                }

                write_to(fd, offset, block, blockSize);
        }

        printf("bfs_fragmenter: marked %lld blocks used\n", occupiedBlocks);

        // write back the superblock

        superBlock.used_blocks += occupiedBlocks;
        write_to(fd, 512, &superBlock, sizeof(superBlock));

        return 0;
}