/**
 * \file
 * \brief VHD management functions (open, close, read write etc)
 */
#ifndef _FILE_OFFSET_BITS
#define _FILE_OFFSET_BITS 64
#endif
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include "minivhd/cwalk.h"
#include "minivhd/libxml2_encoding.h"
#include "minivhd/minivhd_internal.h"
#include "minivhd/minivhd_io.h"
#include "minivhd/minivhd_util.h"
#include "minivhd/minivhd_struct_rw.h"
#include "minivhd/minivhd.h"

int mvhd_errno = 0;
static char tmp_open_path[MVHD_MAX_PATH_BYTES] = {0};
struct MVHDPaths {
        char dir_path[MVHD_MAX_PATH_BYTES];
        char file_name[MVHD_MAX_PATH_BYTES];
        char w2ku_path[MVHD_MAX_PATH_BYTES];
        char w2ru_path[MVHD_MAX_PATH_BYTES];
        char joined_path[MVHD_MAX_PATH_BYTES];
        uint16_t tmp_src_path[MVHD_MAX_PATH_CHARS];
};

static void mvhd_read_footer(MVHDMeta *vhdm);
static void mvhd_read_sparse_header(MVHDMeta *vhdm);
static bool mvhd_footer_checksum_valid(MVHDMeta *vhdm);
static bool mvhd_sparse_checksum_valid(MVHDMeta *vhdm);
static int mvhd_read_bat(MVHDMeta *vhdm, MVHDError *err);
static void mvhd_calc_sparse_values(MVHDMeta *vhdm);
static int mvhd_init_sector_bitmap(MVHDMeta *vhdm, MVHDError *err);

/**
 * \brief Populate data stuctures with content from a VHD footer
 *
 * \param [in] vhdm MiniVHD data structure
 */
static void mvhd_read_footer(MVHDMeta *vhdm) {
        uint8_t buffer[MVHD_FOOTER_SIZE];
        mvhd_fseeko64(vhdm->f, -MVHD_FOOTER_SIZE, SEEK_END);
        fread(buffer, sizeof buffer, 1, vhdm->f);
        mvhd_buffer_to_footer(&vhdm->footer, buffer);
}

/**
 * \brief Populate data stuctures with content from a VHD sparse header
 *
 * \param [in] vhdm MiniVHD data structure
 */
static void mvhd_read_sparse_header(MVHDMeta *vhdm) {
        uint8_t buffer[MVHD_SPARSE_SIZE];
        mvhd_fseeko64(vhdm->f, vhdm->footer.data_offset, SEEK_SET);
        fread(buffer, sizeof buffer, 1, vhdm->f);
        mvhd_buffer_to_header(&vhdm->sparse, buffer);
}

/**
 * \brief Validate VHD footer checksum
 *
 * This works by generating a checksum from the footer, and comparing it against the stored checksum.
 *
 * \param [in] vhdm MiniVHD data structure
 */
static bool mvhd_footer_checksum_valid(MVHDMeta *vhdm) {
        return vhdm->footer.checksum == mvhd_gen_footer_checksum(&vhdm->footer);
}

/**
 * \brief Validate VHD sparse header checksum
 *
 * This works by generating a checksum from the sparse header, and comparing it against the stored checksum.
 *
 * \param [in] vhdm MiniVHD data structure
 */
static bool mvhd_sparse_checksum_valid(MVHDMeta *vhdm) {
        return vhdm->sparse.checksum == mvhd_gen_sparse_checksum(&vhdm->sparse);
}

/**
 * \brief Read BAT into MiniVHD data structure
 *
 * The Block Allocation Table (BAT) is the structure in a sparse and differencing VHD which stores
 * the 4-byte sector offsets for each data block. This function allocates enough memory to contain
 * the entire BAT, and then reads the contents of the BAT into the buffer.
 *
 * \param [in] vhdm MiniVHD data structure
 * \param [out] err this is populated with MVHD_ERR_MEM if the calloc fails
 *
 * \retval -1 if an error occurrs. Check value of err in this case
 * \retval 0 if the function call succeeds
 */
static int mvhd_read_bat(MVHDMeta *vhdm, MVHDError *err) {
        vhdm->block_offset = calloc(vhdm->sparse.max_bat_ent, sizeof *vhdm->block_offset);
        if (vhdm->block_offset == NULL) {
                *err = MVHD_ERR_MEM;
                return -1;
        }
        mvhd_fseeko64(vhdm->f, vhdm->sparse.bat_offset, SEEK_SET);
        uint32_t i;
        for (i = 0; i < vhdm->sparse.max_bat_ent; i++) {
                fread(&vhdm->block_offset[i], sizeof *vhdm->block_offset, 1, vhdm->f);
                vhdm->block_offset[i] = mvhd_from_be32(vhdm->block_offset[i]);
        }
        return 0;
}

/**
 * \brief Perform a one-time calculation of some sparse VHD values
 *
 * \param [in] vhdm MiniVHD data structure
 */
static void mvhd_calc_sparse_values(MVHDMeta *vhdm) {
        vhdm->sect_per_block = vhdm->sparse.block_sz / MVHD_SECTOR_SIZE;
        int bm_bytes = vhdm->sect_per_block / 8;
        vhdm->bitmap.sector_count = bm_bytes / MVHD_SECTOR_SIZE;
        if (bm_bytes % MVHD_SECTOR_SIZE > 0) {
                vhdm->bitmap.sector_count++;
        }
}

/**
 * \brief Allocate memory for a sector bitmap.
 *
 * Each data block is preceded by a sector bitmap. Each bit indicates whether the corresponding sector
 * is considered 'clean' or 'dirty' (for sparse VHD images), or whether to read from the parent or current
 * image (for differencing images).
 *
 * \param [in] vhdm MiniVHD data structure
 * \param [out] err this is populated with MVHD_ERR_MEM if the calloc fails
 *
 * \retval -1 if an error occurrs. Check value of err in this case
 * \retval 0 if the function call succeeds
 */
static int mvhd_init_sector_bitmap(MVHDMeta *vhdm, MVHDError *err) {
        vhdm->bitmap.curr_bitmap = calloc(vhdm->bitmap.sector_count, MVHD_SECTOR_SIZE);
        if (vhdm->bitmap.curr_bitmap == NULL) {
                *err = MVHD_ERR_MEM;
                return -1;
        }
        vhdm->bitmap.curr_block = -1;
        return 0;
}

/**
 * \brief Check if the path for a given platform code exists
 *
 * From the available paths, both relative and absolute, construct a full path
 * and attempt to open a file at that path.
 *
 * Note, this function makes no attempt to verify that the path is the correct
 * VHD image, or even a VHD image at all.
 *
 * \param [in] paths a struct containing all available paths to work with
 * \param [in] the platform code to try and obtain a path for. Setting this to zero
 * will try using the directory of the child image
 *
 * \retval true if a file is found
 * \retval false if a file is not found
 */
static bool mvhd_parent_path_exists(struct MVHDPaths *paths, uint32_t plat_code) {
        memset(paths->joined_path, 0, sizeof paths->joined_path);
        FILE *f;
        int cwk_ret, ferr;
        enum cwk_path_style style = cwk_path_guess_style((const char *)paths->dir_path);
        cwk_path_set_style(style);
        cwk_ret = 1;
        if (plat_code == MVHD_DIF_LOC_W2RU && *paths->w2ru_path) {
                cwk_ret = cwk_path_join((const char *)paths->dir_path, (const char *)paths->w2ru_path, paths->joined_path,
                                        sizeof paths->joined_path);
        } else if (plat_code == MVHD_DIF_LOC_W2KU && *paths->w2ku_path) {
                memcpy(paths->joined_path, paths->w2ku_path, (sizeof paths->joined_path) - 1);
                cwk_ret = 0;
        } else if (plat_code == 0) {
                cwk_ret = cwk_path_join((const char *)paths->dir_path, (const char *)paths->file_name, paths->joined_path,
                                        sizeof paths->joined_path);
        }
        if (cwk_ret > MVHD_MAX_PATH_BYTES) {
                return false;
        }
        f = mvhd_fopen((const char *)paths->joined_path, "rb", &ferr);
        if (f != NULL) {
                /* We found a file at the requested path! */
                memcpy(tmp_open_path, paths->joined_path, (sizeof paths->joined_path) - 1);
                tmp_open_path[sizeof tmp_open_path - 1] = '\0';
                fclose(f);
                return true;
        } else {
                return false;
        }
}

/**
 * \brief attempt to obtain a file path to a file that may be a valid VHD image
 *
 * Differential VHD images store both a UTF-16BE file name (or path), and up to
 * eight "parent locator" entries. Using this information, this function tries to
 * find a parent image.
 *
 * This function does not verify if the path returned is a valid parent image.
 *
 * \param [in] vhdm current MiniVHD data structure
 * \param [out] err any errors that may occurr. Check this if NULL is returned
 *
 * \return a pointer to the global string `tmp_open_path`, or NULL if a path could
 * not be found, or some error occurred
 */
static char *mvhd_get_diff_parent_path(MVHDMeta *vhdm, int *err) {
        int utf_outlen, utf_inlen, utf_ret;
        char *par_fp = NULL;
        /* We can't resolve relative paths if we don't have an absolute
           path to work with */
        if (!cwk_path_is_absolute((const char *)vhdm->filename)) {
                *err = MVHD_ERR_PATH_REL;
                goto end;
        }
        struct MVHDPaths *paths = calloc(1, sizeof *paths);
        if (paths == NULL) {
                *err = MVHD_ERR_MEM;
                goto end;
        }
        size_t dirlen;
        cwk_path_get_dirname((const char *)vhdm->filename, &dirlen);
        if (dirlen >= sizeof paths->dir_path) {
                *err = MVHD_ERR_PATH_LEN;
                goto paths_cleanup;
        }
        memcpy(paths->dir_path, vhdm->filename, dirlen);
        /* Get the filename field from the sparse header. */
        utf_outlen = (int)sizeof paths->file_name;
        utf_inlen = (int)sizeof vhdm->sparse.par_utf16_name;
        utf_ret = UTF16BEToUTF8((unsigned char *)paths->file_name, &utf_outlen,
                                (const unsigned char *)vhdm->sparse.par_utf16_name, &utf_inlen);
        if (utf_ret < 0) {
                mvhd_set_encoding_err(utf_ret, err);
                goto paths_cleanup;
        }
        /* Now read the parent locator entries, both relative and absolute, if they exist */
        unsigned char *loc_path;
        int i;
        for (i = 0; i < 8; i++) {
                utf_outlen = MVHD_MAX_PATH_BYTES - 1;
                if (vhdm->sparse.par_loc_entry[i].plat_code == MVHD_DIF_LOC_W2RU) {
                        loc_path = (unsigned char *)paths->w2ru_path;
                } else if (vhdm->sparse.par_loc_entry[i].plat_code == MVHD_DIF_LOC_W2KU) {
                        loc_path = (unsigned char *)paths->w2ku_path;
                } else {
                        continue;
                }
                utf_inlen = vhdm->sparse.par_loc_entry[i].plat_data_len;
                if (utf_inlen > MVHD_MAX_PATH_BYTES) {
                        *err = MVHD_ERR_PATH_LEN;
                        goto paths_cleanup;
                }
                mvhd_fseeko64(vhdm->f, vhdm->sparse.par_loc_entry[i].plat_data_offset, SEEK_SET);
                fread(paths->tmp_src_path, sizeof(uint8_t), utf_inlen, vhdm->f);
                /* Note, the W2*u parent locators are UTF-16LE, unlike the filename field previously obtained,
                   which is UTF-16BE */
                utf_ret = UTF16LEToUTF8(loc_path, &utf_outlen, (const unsigned char *)paths->tmp_src_path, &utf_inlen);
                if (utf_ret < 0) {
                        mvhd_set_encoding_err(utf_ret, err);
                        goto paths_cleanup;
                }
        }
        /* We have paths in UTF-8. We should have enough info to try and find the parent VHD */
        /* Does the relative path exist? */
        if (mvhd_parent_path_exists(paths, MVHD_DIF_LOC_W2RU)) {
                par_fp = tmp_open_path;
                goto paths_cleanup;
        }
        /* What about trying the child directory? */
        if (mvhd_parent_path_exists(paths, 0)) {
                par_fp = tmp_open_path;
                goto paths_cleanup;
        }
        /* Well, all else fails, try the stored absolute path, if it exists */
        if (mvhd_parent_path_exists(paths, MVHD_DIF_LOC_W2KU)) {
                par_fp = tmp_open_path;
                goto paths_cleanup;
        }
        /* If we reach this point, we could not find a path with a valid file */
        par_fp = NULL;
        *err = MVHD_ERR_PAR_NOT_FOUND;

paths_cleanup:
        free(paths);
        paths = NULL;
end:
        return par_fp;
}

/**
 * \brief Attach the read/write function pointers to read/write functions
 *
 * Depending on the VHD type, different sector reading and writing functions are used.
 * The functions are called via function pointers stored in the vhdm struct.
 *
 * \param [in] vhdm MiniVHD data structure
 */
static void mvhd_assign_io_funcs(MVHDMeta *vhdm) {
        switch (vhdm->footer.disk_type) {
        case MVHD_TYPE_FIXED:
                vhdm->read_sectors = mvhd_fixed_read;
                vhdm->write_sectors = mvhd_fixed_write;
                break;
        case MVHD_TYPE_DYNAMIC:
                vhdm->read_sectors = mvhd_sparse_read;
                vhdm->write_sectors = mvhd_sparse_diff_write;
                break;
        case MVHD_TYPE_DIFF:
                vhdm->read_sectors = mvhd_diff_read;
                vhdm->write_sectors = mvhd_sparse_diff_write;
                break;
        }
        if (vhdm->readonly) {
                vhdm->write_sectors = mvhd_noop_write;
        }
}

bool mvhd_file_is_vhd(FILE *f) {
        if (f) {
                uint8_t con_str[8];
                mvhd_fseeko64(f, -MVHD_FOOTER_SIZE, SEEK_END);
                fread(con_str, sizeof con_str, 1, f);
                return mvhd_is_conectix_str(con_str);
        } else {
                return false;
        }
}

MVHDGeom mvhd_calculate_geometry(uint64_t size) {
        MVHDGeom chs;
        uint32_t ts = (uint32_t)(size / MVHD_SECTOR_SIZE);
        uint32_t spt, heads, cyl, cth;
        if (ts > 65535 * 16 * 255) {
                ts = 65535 * 16 * 255;
        }
        if (ts >= 65535 * 16 * 63) {
                spt = 255;
                heads = 16;
                cth = ts / spt;
        } else {
                spt = 17;
                cth = ts / spt;
                heads = (cth + 1023) / 1024;
                if (heads < 4) {
                        heads = 4;
                }
                if (cth >= (heads * 1024) || heads > 16) {
                        spt = 31;
                        heads = 16;
                        cth = ts / spt;
                }
                if (cth >= (heads * 1024)) {
                        spt = 63;
                        heads = 16;
                        cth = ts / spt;
                }
        }
        cyl = cth / heads;
        chs.heads = heads;
        chs.spt = spt;
        chs.cyl = cyl;
        return chs;
}

MVHDMeta *mvhd_open(const char *path, bool readonly, int *err) {
        MVHDError open_err;
        MVHDMeta *vhdm = calloc(sizeof *vhdm, 1);
        if (vhdm == NULL) {
                *err = MVHD_ERR_MEM;
                goto end;
        }
        if (strlen(path) >= sizeof vhdm->filename) {
                *err = MVHD_ERR_PATH_LEN;
                goto cleanup_vhdm;
        }
        // This is safe, as we've just checked for potential overflow above
        strcpy(vhdm->filename, path);
        vhdm->f = readonly ? mvhd_fopen((const char *)vhdm->filename, "rb", err)
                           : mvhd_fopen((const char *)vhdm->filename, "rb+", err);
        if (vhdm->f == NULL) {
                /* note, mvhd_fopen sets err for us */
                goto cleanup_vhdm;
        }
        vhdm->readonly = readonly;
        if (!mvhd_file_is_vhd(vhdm->f)) {
                *err = MVHD_ERR_NOT_VHD;
                goto cleanup_file;
        }
        mvhd_read_footer(vhdm);
        if (!mvhd_footer_checksum_valid(vhdm)) {
                *err = MVHD_ERR_FOOTER_CHECKSUM;
                goto cleanup_file;
        }
        if (vhdm->footer.disk_type == MVHD_TYPE_DIFF || vhdm->footer.disk_type == MVHD_TYPE_DYNAMIC) {
                mvhd_read_sparse_header(vhdm);
                if (!mvhd_sparse_checksum_valid(vhdm)) {
                        *err = MVHD_ERR_SPARSE_CHECKSUM;
                        goto cleanup_file;
                }
                if (mvhd_read_bat(vhdm, &open_err) == -1) {
                        *err = open_err;
                        goto cleanup_file;
                }
                mvhd_calc_sparse_values(vhdm);
                if (mvhd_init_sector_bitmap(vhdm, &open_err) == -1) {
                        *err = open_err;
                        goto cleanup_bat;
                }

        } else if (vhdm->footer.disk_type != MVHD_TYPE_FIXED) {
                *err = MVHD_ERR_TYPE;
                goto cleanup_bitmap;
        }
        mvhd_assign_io_funcs(vhdm);
        vhdm->format_buffer.zero_data = calloc(64, MVHD_SECTOR_SIZE);
        if (vhdm->format_buffer.zero_data == NULL) {
                *err = MVHD_ERR_MEM;
                goto cleanup_bitmap;
        }
        vhdm->format_buffer.sector_count = 64;
        if (vhdm->footer.disk_type == MVHD_TYPE_DIFF) {
                char *par_path = mvhd_get_diff_parent_path(vhdm, err);
                if (par_path == NULL) {
                        goto cleanup_format_buff;
                }
                uint32_t par_mod_ts = mvhd_file_mod_timestamp(par_path, err);
                if (*err != 0) {
                        goto cleanup_format_buff;
                }
                if (vhdm->sparse.par_timestamp != par_mod_ts) {
                        /* The last-modified timestamp is to fragile to make this a fatal error.
                           Instead, we inform the caller of the potential problem. */
                        *err = MVHD_ERR_TIMESTAMP;
                }
                vhdm->parent = mvhd_open(par_path, true, err);
                if (vhdm->parent == NULL) {
                        goto cleanup_format_buff;
                }
                if (memcmp(vhdm->sparse.par_uuid, vhdm->parent->footer.uuid, sizeof vhdm->sparse.par_uuid) != 0) {
                        *err = MVHD_ERR_INVALID_PAR_UUID;
                        goto cleanup_format_buff;
                }
        }
        /* If we've reached this point, we are good to go, so skip the cleanup steps */
        goto end;
cleanup_format_buff:
        free(vhdm->format_buffer.zero_data);
        vhdm->format_buffer.zero_data = NULL;
cleanup_bitmap:
        free(vhdm->bitmap.curr_bitmap);
        vhdm->bitmap.curr_bitmap = NULL;
cleanup_bat:
        free(vhdm->block_offset);
        vhdm->block_offset = NULL;
cleanup_file:
        fclose(vhdm->f);
        vhdm->f = NULL;
cleanup_vhdm:
        free(vhdm);
        vhdm = NULL;
end:
        return vhdm;
}

void mvhd_close(MVHDMeta *vhdm) {
        if (vhdm != NULL) {
                if (vhdm->parent != NULL) {
                        mvhd_close(vhdm->parent);
                }
                fclose(vhdm->f);
                if (vhdm->block_offset != NULL) {
                        free(vhdm->block_offset);
                        vhdm->block_offset = NULL;
                }
                if (vhdm->bitmap.curr_bitmap != NULL) {
                        free(vhdm->bitmap.curr_bitmap);
                        vhdm->bitmap.curr_bitmap = NULL;
                }
                if (vhdm->format_buffer.zero_data != NULL) {
                        free(vhdm->format_buffer.zero_data);
                        vhdm->format_buffer.zero_data = NULL;
                }
                free(vhdm);
                vhdm = NULL;
        }
}

int mvhd_diff_update_par_timestamp(MVHDMeta *vhdm, int *err) {
        uint8_t sparse_buff[1024];
        if (vhdm == NULL || err == NULL) {
                *err = MVHD_ERR_INVALID_PARAMS;
                return -1;
        }
        if (vhdm->footer.disk_type != MVHD_TYPE_DIFF) {
                *err = MVHD_ERR_TYPE;
                return -1;
        }
        char *par_path = mvhd_get_diff_parent_path(vhdm, err);
        if (par_path == NULL) {
                return -1;
        }
        uint32_t par_mod_ts = mvhd_file_mod_timestamp(par_path, err);
        if (*err != 0) {
                return -1;
        }
        /* Update the timestamp and sparse header checksum */
        vhdm->sparse.par_timestamp = par_mod_ts;
        vhdm->sparse.checksum = mvhd_gen_sparse_checksum(&vhdm->sparse);
        /* Generate and write the updated sparse header */
        mvhd_header_to_buffer(&vhdm->sparse, sparse_buff);
        mvhd_fseeko64(vhdm->f, (int64_t)vhdm->footer.data_offset, SEEK_SET);
        fwrite(sparse_buff, sizeof sparse_buff, 1, vhdm->f);
        return 0;
}

int mvhd_read_sectors(MVHDMeta *vhdm, uint32_t offset, int num_sectors, void *out_buff) {
        return vhdm->read_sectors(vhdm, offset, num_sectors, out_buff);
}

int mvhd_write_sectors(MVHDMeta *vhdm, uint32_t offset, int num_sectors, void *in_buff) {
        return vhdm->write_sectors(vhdm, offset, num_sectors, in_buff);
}

int mvhd_format_sectors(MVHDMeta *vhdm, uint32_t offset, int num_sectors) {
        int num_full = num_sectors / vhdm->format_buffer.sector_count;
        int remain = num_sectors % vhdm->format_buffer.sector_count;
        int i;
        for (i = 0; i < num_full; i++) {
                vhdm->write_sectors(vhdm, offset, vhdm->format_buffer.sector_count, vhdm->format_buffer.zero_data);
                offset += vhdm->format_buffer.sector_count;
        }
        vhdm->write_sectors(vhdm, offset, remain, vhdm->format_buffer.zero_data);
        return 0;
}
