#include <stdint.h>
#include <string.h>
#include "usecfs.h"
#include "usecfs_dev.h"


#ifndef NULL
#   define NULL (void *)0
#endif

#define NO_BLOCK (0xFFFFFFFFUL)
#ifndef BLOCKDEV_OPEN_ARGS
#  define BLOCKDEV_OPEN_ARGS (NULL)
#endif

#define MAX_BLOCKS_0 ((BLOCK_SIZE - MAX_FILENAME) / 4) - 3 /* accounts for 3 32bit fields + FILENAME */
#define MAX_BLOCKS_N ((BLOCK_SIZE - MAX_FILENAME) / 4) - 1 /* overhead is only 'extra' fields        */
#define MAX_INLINE_SIZE ((MAX_BLOCKS_0 - 1) * 4)           /* Max data size for self-contained block */
#define INLINE_PAYLOAD(x) ((uint8_t *)(&x->blk[1]))        /* Macro to access INLINE paylaod         */

#include <wolfssl/wolfcrypt/settings.h>
#include <wolfssl/options.h>
#include <wolfssl/wolfcrypt/sha256.h>
#ifdef CRYPTO
#define CRYPTO_BLOCK_SIZE 16

uint8_t crypto_tmp[CRYPTO_BLOCK_SIZE];
uint8_t crypto_iv[CRYPTO_BLOCK_SIZE];

#include <wolfssl/wolfcrypt/chacha.h>
#include <wolfssl/wolfcrypt/pwdbased.h>

static ChaCha chacha;

ChaCha *chacha_secret = &chacha;

#define CRYPTO_KEY_SIZE 32
#endif

#define HASH_LEN 32
#define MAGIC 0x5AFED15C

static uint8_t cache[BLOCK_SIZE];
static uint8_t inline_buffer_copy[MAX_INLINE_SIZE];


struct __attribute__((packed)) extra {
    uint32_t extra;
    uint32_t blk[(BLOCK_SIZE / 4) - 1];
};

struct __attribute__((packed)) inode {
    uint32_t extra;
    uint32_t blk[((BLOCK_SIZE - MAX_FILENAME) / 4) - 3];
    uint32_t size;
    uint32_t nextfile;
    char filename[MAX_FILENAME];
};

struct __attribute__((packed)) root_block {
    uint32_t magic;
    uint32_t blk[((BLOCK_SIZE - MAX_FILENAME) / 4) - 3];
    uint32_t fs_size; /* unused */
    uint32_t nextfile;
    uint8_t uuid[UUID_LEN];
    uint8_t hash[HASH_LEN];
};


/* One-sector cache, single entry point for read/write
 * on blocks
 */
static uint32_t cached_block = NO_BLOCK;
void *blockdev = NULL;

static int is_block_empty(void)
{
    uint32_t *w = (uint32_t *)cache;
    int i;
    for(i = 0; i < (BLOCK_SIZE / sizeof(uint32_t)); i++) {
        if (w[i] != NO_BLOCK)
            return 0;
    }
    return 1;
}

static void cache_commit(void)
{
    if (cached_block == NO_BLOCK)
        return;
#ifdef CRYPTO
        if (!is_block_empty()) {
            uint32_t i;
            memset(crypto_iv, 0, CRYPTO_BLOCK_SIZE);
            for (i = 0; i < BLOCK_SIZE / CRYPTO_BLOCK_SIZE; i++) {
                memcpy(&crypto_iv[0], &cached_block, sizeof(uint32_t));
                memcpy(&crypto_iv[4], &i, sizeof(uint32_t));
                memcpy(crypto_tmp, cache + (i * CRYPTO_BLOCK_SIZE), CRYPTO_BLOCK_SIZE);
                wc_Chacha_SetIV(&chacha, crypto_iv, CRYPTO_BLOCK_SIZE);
                wc_Chacha_Process(&chacha, cache + i * CRYPTO_BLOCK_SIZE, crypto_tmp, CRYPTO_BLOCK_SIZE);
            }
        }
#endif
    block_write(blockdev, cache, cached_block);
    cached_block = NO_BLOCK;
}

static void cache_load(uint32_t blk)
{
    if (cached_block == blk)
        return;
    if (cached_block != NO_BLOCK)
        cache_commit();
    if (block_read(blockdev, cache, blk) == BLOCK_SIZE) {
        cached_block = blk;
#ifdef CRYPTO
        if (!is_block_empty()) {
            uint32_t i;
            memset(crypto_iv, 0, CRYPTO_BLOCK_SIZE);
            for (i = 0; i < BLOCK_SIZE / CRYPTO_BLOCK_SIZE; i++) {
                memcpy(&crypto_iv[0], &blk, sizeof(uint32_t));
                memcpy(&crypto_iv[4], &i, sizeof(uint32_t));
                memcpy(crypto_tmp, cache + (i * CRYPTO_BLOCK_SIZE), CRYPTO_BLOCK_SIZE);
                wc_Chacha_SetIV(&chacha, crypto_iv, CRYPTO_BLOCK_SIZE);
                wc_Chacha_Process(&chacha, cache + i * CRYPTO_BLOCK_SIZE, crypto_tmp, CRYPTO_BLOCK_SIZE);
            }
        }
#endif
    }
}

static uint32_t get_file(const char *filename)
{
    struct inode *ci = (struct inode *)cache;
    uint32_t blk = 0;
    cache_load(blk);
    while (1) {
        if (strncmp(filename, ci->filename, MAX_FILENAME - 1) == 0)
            return blk;
        if (ci->nextfile != NO_BLOCK) {
            blk = ci->nextfile;
            cache_load(blk);
        } else
            return NO_BLOCK;
    }
    return NO_BLOCK; /* Never reached */
}

static uint32_t get_free_block(uint32_t *fs_tail)
{
    struct inode *ci = (struct inode *)cache;
    struct extra *ce = (struct extra *)cache;
    uint32_t first_free = 1;
    uint32_t cur_inode_0 = 0;
    int i;
    cache_load(0);
    if (ci->nextfile == NO_BLOCK) {
        if (fs_tail)
            *fs_tail = 0;
        return 1;
    }
    cur_inode_0 = ci->nextfile;
    cache_load(cur_inode_0);
    while (1) {
        for (i = 0; i < MAX_BLOCKS_0; i++)
        {
            if (ci->blk[i] == NO_BLOCK)
                break;
            if (first_free == ci->blk[i])
                first_free++;
        }
        if (ci->extra != NO_BLOCK)
        {
            cache_load(ci->extra);
            while (1) {
                for (i = 0; i < MAX_BLOCKS_N; i++)
                {
                    if (ce->blk[i] == NO_BLOCK)
                        break;
                    if (first_free == ce->blk[i])
                        first_free++;
                }
                if (ce->extra == NO_BLOCK)
                    break;
                cache_load(ce->extra);
            }
            cache_load(cur_inode_0);
        }
        if (ci->nextfile != NO_BLOCK) {
            cur_inode_0 = ci->nextfile;
            cache_load(cur_inode_0);
        } else {
            if (fs_tail)
                *fs_tail = cur_inode_0;
            return first_free;
        }
    }
}

static uint32_t new_inode(void)
{
    struct inode *ci = (struct inode *)cache;
    uint32_t fs_tail = NO_BLOCK;
    uint32_t new_block = get_free_block(&fs_tail);
    if (new_block == NO_BLOCK || fs_tail == NO_BLOCK)
        return NO_BLOCK;

    cache_load(fs_tail);
    ci->nextfile = new_block;
    cache_load(new_block);
    memset(cache, 0xFF, BLOCK_SIZE);
    ci->nextfile = NO_BLOCK;
    return new_block;
}

struct openfile {
    uint32_t blk;
    uint32_t off;
};

struct openfile OpenFiles[MAX_OPEN_FILES];
static int get_free_fd(void)
{
    int i, free_fd = -1;

    for (i = 0; i < MAX_OPEN_FILES; i++) {
        if (OpenFiles[i].blk == NO_BLOCK) {
            free_fd = i;
            break;
        }
    }
    return free_fd;
}

static uint32_t get_index_block(uint32_t inode0, uint32_t off)
{
    struct inode *ci = (struct inode *)cache;
    struct extra *ce = (struct extra *)cache;
    uint32_t idx = 0;
    cache_load(inode0);
    idx = off / BLOCK_SIZE;
    if (idx <= MAX_BLOCKS_0)
        return inode0;
    if (ci->extra == NO_BLOCK)
        return NO_BLOCK;
    inode0 = ci->extra;
    cache_load(inode0);
    idx -= MAX_BLOCKS_0;
    while (idx >= MAX_BLOCKS_N) {
        if (ce->extra == NO_BLOCK)
            return NO_BLOCK;
        inode0 = ce->extra;
        cache_load(inode0);
        idx -= MAX_BLOCKS_N;
    }
    return inode0;
}

static void set_block(uint32_t node0, uint32_t nodeN, uint32_t off)
{
    struct inode *ci = (struct inode *)cache;
    struct extra *ce = (struct extra *)cache;
    uint32_t idx;
    uint32_t idxblk;
    idxblk = get_index_block(node0, off);
}

static void file_grow(uint32_t node0, uint32_t newsize)
{
    struct inode *ci = (struct inode *)cache;
    struct extra *ce = (struct extra *)cache;
    uint32_t cur_idx, new_idx;
    uint32_t cur_idx_block, new_idx_block;
    uint32_t i;
    cache_load(node0);
    cur_idx = ci->size / BLOCK_SIZE;
    new_idx = newsize / BLOCK_SIZE;
    cur_idx_block = get_index_block(node0, ci->size);
    new_idx_block = get_index_block(node0, newsize);

    for (i = cur_idx; i < new_idx; i++) {
        cur_idx_block = get_index_block(node0, i * BLOCK_SIZE);
        if (cur_idx_block == NO_BLOCK)
            return;
        cache_load(cur_idx_block);
        if (cur_idx_block == node0) {
            uint32_t idx = i + 1;
            if (idx == (MAX_BLOCKS_0)) {
                uint32_t extra = get_free_block(NULL);
                cache_load(extra);
                memset(cache, 0xFF, BLOCK_SIZE);
                cache_load(node0);
                ci->extra = extra;
                cache_load(extra);
                idx = 0;
            }
            ci->blk[idx] = get_free_block(NULL);
        } else {
            uint32_t idx = ((i - MAX_BLOCKS_0) % MAX_BLOCKS_N) + 1;
            if (idx == MAX_BLOCKS_N) {
                uint32_t extra = get_free_block(NULL);
                cache_load(extra);
                memset(cache, 0xFF, BLOCK_SIZE);
                cache_load(cur_idx_block);
                ce->extra = extra;
                cache_load(extra);
                idx = 0;
            }
            ce->blk[idx] = get_free_block(NULL);
        }
    }
    if (ci->size < newsize)
        ci->size = newsize;
}

/* Public interface */

int usecfs_format(const uint8_t *uuid)
{
    struct root_block *rb = (struct root_block *)cache;
    wc_Sha256 sha;
    cache_load(0);
    memset(cache, 0xFF, BLOCK_SIZE);
    rb->magic = MAGIC;
    rb->blk[0] = 0;
    memcpy(rb->uuid, uuid, UUID_LEN);
    wc_InitSha256(&sha);
    wc_Sha256Update(&sha, uuid, UUID_LEN);
    wc_Sha256Final(&sha, rb->hash);
    cache_commit();
    return 0;
}

int usecfs_mount(uint8_t *uuid)
{
    struct root_block *rb = (struct root_block *)cache;
    wc_Sha256 sha;
    uint8_t hash[HASH_LEN];
    cache_load(0);
    if (rb->magic != MAGIC)
        return -1;
    wc_InitSha256(&sha);
    wc_Sha256Update(&sha, rb->uuid, UUID_LEN);
    wc_Sha256Final(&sha, hash);
    if (memcmp(hash, rb->hash, HASH_LEN) != 0)
        return -1;
    if (uuid)
        memcpy(uuid, rb->uuid, UUID_LEN);
    return 0;
}

int usecfs_read(int fd, void *data, uint32_t len)
{
    int r = 0;
    uint32_t node0 = OpenFiles[fd].blk;
    uint32_t *off = &OpenFiles[fd].off;
    struct inode *ci = (struct inode *)cache;
    if (node0 == NO_BLOCK)
        return -1;
    cache_load(node0);
    if ((ci->size - *off) < len)
        len = ci->size - *off;
    if (ci->blk[0] == node0) {
        /* file is INLINE */
        memcpy(data, INLINE_PAYLOAD(ci), len);
        *off += len;
        return len;
    }
    while (r < len) {
        int idx = *off / BLOCK_SIZE;
        uint32_t idx_block;
        uint32_t off_within_block;
        uint32_t len_within_block;

        /* Find the index block for the offset */
        idx_block = get_index_block(node0, *off);
        if (idx_block == NO_BLOCK)
            return -1;
        cache_load(idx_block);
        if (ci->blk[idx] == NO_BLOCK)
            return -1;
        cache_load(ci->blk[idx]);

        /* Find the offset/len within this block */
        off_within_block = *off % BLOCK_SIZE;
        len_within_block = BLOCK_SIZE - off_within_block;
        if (len_within_block > (len - r))
            len_within_block = (len - r);

        /* Copy data from cache, increase counters */
        memcpy(data + r, cache + off_within_block, len_within_block);
        r += len_within_block;
        *off += len_within_block;
    }
    return len;
}

int usecfs_write(int fd, const void *data, uint32_t len)
{
    int w = 0;
    uint32_t node0 = OpenFiles[fd].blk;
    uint32_t *off = &OpenFiles[fd].off;
    struct inode *ci = (struct inode *)cache;
    uint32_t blkn;
    if (node0 == NO_BLOCK)
        return -1;
    cache_load(node0);
    if (ci->size + *off <= MAX_INLINE_SIZE) {
        if (ci->size + *off + len <= MAX_INLINE_SIZE) {
            /* file can still be inline */
            memcpy(INLINE_PAYLOAD(ci) + *off, data, len);
            *off += len;
            if (*off > ci->size)
                ci->size = *off;
            cache_commit();
            return len;
        } else {
            /* Special case: migrating data to new block */
            memcpy(inline_buffer_copy, INLINE_PAYLOAD(ci), MAX_INLINE_SIZE);
            blkn = get_free_block(NULL);
            memcpy(cache, inline_buffer_copy, MAX_INLINE_SIZE);
            cache_load(node0);
            ci->blk[0] = blkn;
        }
    }
    if (ci->size < *off + len)
        file_grow(node0, *off + len);

    while (w < len) {
        uint32_t off_within_block;
        uint32_t len_within_block;
        int idx = *off / BLOCK_SIZE;
        uint32_t idx_block;

        idx_block = get_index_block(node0, *off);
        if (idx_block == NO_BLOCK)
            return -1;

        cache_load(idx_block);
        off_within_block = *off % BLOCK_SIZE;
        len_within_block = BLOCK_SIZE - off_within_block;
        if (len_within_block > (len - w))
            len_within_block = (len - w);
        if (ci->blk[idx] == NO_BLOCK)
            return -1;
        cache_load(ci->blk[idx]);
        memcpy(cache + off_within_block, data + w, len_within_block);
        w += len_within_block;
        *off += len_within_block;
    }
    return len;
}

int usecfs_seek(int fd, int offset, int whence)
{
    struct inode *ci = (struct inode *)cache;
    if (OpenFiles[fd].blk == NO_BLOCK)
        return -1;
    cache_load(OpenFiles[fd].blk);
    switch(whence) {
        case 0: /* SEEK_SET */
            if (offset < 0)
                return -1;
            OpenFiles[fd].off = offset;
            break;
        case 1: /* SEEK_CUR */
            if ((OpenFiles[fd].off - offset) < 0)
                OpenFiles[fd].off = 0;
            else
                OpenFiles[fd].off += offset;
            break;
        case 2: /* SEEK_END */
            OpenFiles[fd].off = ci->size + offset;
            break;
        default:
            return -1;
    } /* switch(whence) */
    if (OpenFiles[fd].off > ci->size)
        file_grow(OpenFiles[fd].blk, OpenFiles[fd].off);
    return OpenFiles[fd].off;
}

int usecfs_open(const char *name)
{
    int free_fd;
    free_fd = get_free_fd();
    if (free_fd < 0)
        return -1;
    OpenFiles[free_fd].blk = get_file(name);
    if (OpenFiles[free_fd].blk == NO_BLOCK)
        return -1;
    OpenFiles[free_fd].off = 0;
    return free_fd;
}

int usecfs_creat(const char *name)
{
    int free_fd;
    int i;
    uint32_t newnode;
    struct inode *ci = (struct inode *)cache;
    free_fd = get_free_fd();
    if (free_fd < 0)
        return -1;
    if (get_file(name) != NO_BLOCK)
        return -1;
    newnode = new_inode();
    if (newnode == NO_BLOCK)
        return -1;
    OpenFiles[free_fd].blk = newnode;

    /* Initialize inode (already in cache from new_inode()) */
    ci->size = 0; /* newly created file is zero bytes long */
    memset(ci->filename, 0, MAX_FILENAME);
    for (i = 0; (i < strlen(name)) && (i < MAX_FILENAME - 1); i++)
        ci->filename[i] = name[i];
    ci->blk[0] = newnode; /* self-reference, start as INLINE */
    for (i = 1; i < MAX_BLOCKS_0; i++)
        ci->blk[i] = NO_BLOCK;
    OpenFiles[free_fd].off = 0;
    cache_commit();
    return free_fd;
}

int usecfs_truncate(int fd, uint32_t newsize)
{
    struct inode *ci = (struct inode *)cache;
    struct extra *ce = (struct extra *)cache;
    uint32_t idx, new_idx;
    uint32_t last_block_idx;
    uint32_t idx_block, new_idx_block;
    uint32_t node0 = OpenFiles[fd].blk;
    uint32_t idx_off;
    uint32_t i;
    if (node0 == NO_BLOCK)
        return -1;
    cache_load(OpenFiles[fd].blk);
    if (ci->size <= newsize)
        return -1;
    idx = ci->size / BLOCK_SIZE;
    new_idx = newsize / BLOCK_SIZE;
    if (ci->size <= (BLOCK_SIZE * MAX_BLOCKS_0) || /* File is inlined, or            */
         (idx == new_idx) )                       /* number of blocks is unchanged? */
    {
        ci->size = newsize;
        return 0;
    }
    if ((ci->size > MAX_INLINE_SIZE) && (newsize <= MAX_INLINE_SIZE))
    {
        /* Special case: file was big, now it's inlined. */
        cache_load(ci->blk[0]);
        memcpy(inline_buffer_copy, cache, newsize);
        cache_load(node0);
        memcpy(INLINE_PAYLOAD(ci), inline_buffer_copy, newsize);
        ci->blk[0] = node0; /* establish self-reference */
        for (i = 1; i < MAX_BLOCKS_0; i++)
            ci->blk[i] = NO_BLOCK;
        ci->size = newsize;
        return 0;
    }
    idx_block = get_index_block(node0, ci->size);
    new_idx_block = get_index_block(node0, newsize);
    cache_load(new_idx_block);
    if (idx_block != new_idx_block)
        ci->extra = NO_BLOCK;
    if (idx_block <= MAX_BLOCKS_0) {
        idx_off = new_idx_block;
        for (i = idx_off; i < MAX_BLOCKS_N; i++)
           ce->blk[i] = NO_BLOCK; 
    } else {
        idx_off = (new_idx_block - MAX_BLOCKS_0) % MAX_BLOCKS_N;
        while (idx_off > MAX_BLOCKS_N)
            idx_off -= MAX_BLOCKS_N;
        for (i = idx_off; i < MAX_BLOCKS_N; i++)
           ce->blk[i] = NO_BLOCK; 
    }
    ci->size = newsize;
    return 0;
}

int usecfs_unlink(const char *filename)
{
    struct inode *ci = (struct inode *)cache;
    uint32_t blk = 1;
    uint32_t blk_prev = 0;
    uint32_t blk_next;
    cache_load(blk);
    while (1) {
        if (strncmp(filename, ci->filename, MAX_FILENAME - 1) == 0)
        {
            blk_next = ci->nextfile;
            cache_load(blk_prev);
            ci->nextfile = blk_next;
            return 0;
        }
        if (ci->nextfile != NO_BLOCK) {
            blk_prev = blk;
            blk = ci->nextfile;
            cache_load(blk);
        } else
            return -1;
    }
    return -1; /* Never reached */
}

int usecfs_close(int fd)
{
    if (OpenFiles[fd].blk != NO_BLOCK) {
        cache_commit();
        OpenFiles[fd].blk = NO_BLOCK;
    }
    return 0;
}

#define SALT_LEN 32
const uint8_t password_salt[SALT_LEN] = {
  0xe7, 0xa1, 0x9c, 0xb0, 0x48, 0xa8, 0x30, 0xf9, 0x37, 0xda, 0x8e, 0xde,
  0xff, 0xb2, 0x62, 0x03, 0x24, 0x55, 0xb8, 0x8b, 0x7b, 0x18, 0x68, 0x57,
  0x7d, 0x35, 0xbe, 0xbd, 0xf6, 0x0e, 0xc1, 0x2c
};

static uint8_t chacha_key[CRYPTO_KEY_SIZE];
int usecfs_getkey(void)
{
    return chacha_key;
}

int usecfs_init(const char *password, int format, uint8_t *uuid)
{
    blockdev = block_open(BLOCKDEV_OPEN_ARGS);
    if (!blockdev)
        return -1;
    memset(OpenFiles, 0xFF, MAX_OPEN_FILES * sizeof(struct openfile));
    {
        int ret = 0;
        ret = wc_PBKDF2(chacha_key, password, strlen(password), password_salt, SALT_LEN, 2048, CRYPTO_KEY_SIZE, SHA256);
        wc_Chacha_SetKey(&chacha, chacha_key, CRYPTO_KEY_SIZE);
    }
    if (format) {
        if (!uuid)
            return -1;
        else
            return usecfs_format(uuid);
    } else
        return usecfs_mount(uuid);
}