#ifndef _FILE_OFFSET_BITS
#define _FILE_OFFSET_BITS 64
#endif
#include <assert.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include "minivhd/cwalk.h"
#include "minivhd/libxml2_encoding.h"
#include "minivhd/minivhd_internal.h"
#include "minivhd/minivhd_util.h"
#include "minivhd/minivhd_struct_rw.h"
#include "minivhd/minivhd_io.h"
#include "minivhd/minivhd_create.h"
#include "minivhd/minivhd.h"

static void mvhd_gen_footer(MVHDFooter *footer, uint64_t size_in_bytes, MVHDGeom *geom, MVHDType type,
                            uint64_t sparse_header_off);
static void mvhd_gen_sparse_header(MVHDSparseHeader *header, uint32_t num_blks, uint64_t bat_offset,
                                   uint32_t block_size_in_sectors);
static int mvhd_gen_par_loc(MVHDSparseHeader *header, const char *child_path, const char *par_path, uint64_t start_offset,
                            mvhd_utf16 *w2ku_path_buff, mvhd_utf16 *w2ru_path_buff, MVHDError *err);
static MVHDMeta *mvhd_create_sparse_diff(const char *path, const char *par_path, uint64_t size_in_bytes, MVHDGeom *geom,
                                         uint32_t block_size_in_sectors, int *err);

/**
 * \brief Populate a VHD footer
 *
 * \param [in] footer to populate
 * \param [in] size_in_bytes is the total size of the virtual hard disk in bytes
 * \param [in] geom to use
 * \param [in] type of HVD that is being created
 * \param [in] sparse_header_off, an absolute file offset to the sparse header. Not used for fixed VHD images
 */
static void mvhd_gen_footer(MVHDFooter *footer, uint64_t size_in_bytes, MVHDGeom *geom, MVHDType type,
                            uint64_t sparse_header_off) {
        memcpy(footer->cookie, "conectix", sizeof footer->cookie);
        footer->features = 0x00000002;
        footer->fi_fmt_vers = 0x00010000;
        footer->data_offset = (type == MVHD_TYPE_DIFF || type == MVHD_TYPE_DYNAMIC) ? sparse_header_off : 0xffffffffffffffff;
        footer->timestamp = vhd_calc_timestamp();
        memcpy(footer->cr_app, "mvhd", sizeof footer->cr_app);
        footer->cr_vers = 0x000e0000;
        memcpy(footer->cr_host_os, "Wi2k", sizeof footer->cr_host_os);
        footer->orig_sz = footer->curr_sz = size_in_bytes;
        footer->geom.cyl = geom->cyl;
        footer->geom.heads = geom->heads;
        footer->geom.spt = geom->spt;
        footer->disk_type = type;
        mvhd_generate_uuid(footer->uuid);
        footer->checksum = mvhd_gen_footer_checksum(footer);
}

/**
 * \brief Populate a VHD sparse header
 *
 * \param [in] header for sparse and differencing images
 * \param [in] num_blks is the number of data blocks that the image contains
 * \param [in] bat_offset is the absolute file offset for start of the Block Allocation Table
 * \param [in] block_size_in_sectors is the block size in sectors.
 */
static void mvhd_gen_sparse_header(MVHDSparseHeader *header, uint32_t num_blks, uint64_t bat_offset,
                                   uint32_t block_size_in_sectors) {
        memcpy(header->cookie, "cxsparse", sizeof header->cookie);
        header->data_offset = 0xffffffffffffffff;
        header->bat_offset = bat_offset;
        header->head_vers = 0x00010000;
        header->max_bat_ent = num_blks;
        header->block_sz = block_size_in_sectors * (uint32_t)MVHD_SECTOR_SIZE;
        header->checksum = mvhd_gen_sparse_checksum(header);
}

/**
 * \brief Generate parent locators for differencing VHD images
 *
 * \param [in] header the sparse header to populate with parent locator entries
 * \param [in] child_path is the full path to the VHD being created
 * \param [in] par_path is the full path to the parent image
 * \param [in] start_offset is the absolute file offset from where to start storing the entries themselves. Must be sector
 * aligned. \param [out] w2ku_path_buff is a buffer containing the full path to the parent, encoded as UTF16-LE \param [out]
 * w2ru_path_buff is a buffer containing the relative path to the parent, encoded as UTF16-LE \param [out] err indicates what
 * error occurred, if any
 *
 * \retval 0 if success
 * \retval < 0 if an error occurrs. Check value of *err for actual error
 */
static int mvhd_gen_par_loc(MVHDSparseHeader *header, const char *child_path, const char *par_path, uint64_t start_offset,
                            mvhd_utf16 *w2ku_path_buff, mvhd_utf16 *w2ru_path_buff, MVHDError *err) {
        /* Get our paths to store in the differencing VHD. We want both the absolute path to the parent,
           as well as the relative path from the child VHD */
        int rv = 0;
        char *par_filename;
        size_t par_fn_len;
        char rel_path[MVHD_MAX_PATH_BYTES] = {0};
        char child_dir[MVHD_MAX_PATH_BYTES] = {0};
        size_t child_dir_len;
        if (strlen(child_path) < sizeof child_dir) {
                strcpy(child_dir, child_path);
        } else {
                *err = MVHD_ERR_PATH_LEN;
                rv = -1;
                goto end;
        }
        cwk_path_get_basename(par_path, (const char **)&par_filename, &par_fn_len);
        cwk_path_get_dirname(child_dir, &child_dir_len);
        child_dir[child_dir_len] = '\0';
        size_t rel_len = cwk_path_get_relative(child_dir, par_path, rel_path, sizeof rel_path);
        if (rel_len > sizeof rel_path) {
                *err = MVHD_ERR_PATH_LEN;
                rv = -1;
                goto end;
        }
        /* We have our paths, now store the parent filename directly in the sparse header. */
        int outlen = sizeof header->par_utf16_name;
        int utf_ret;
        utf_ret = UTF8ToUTF16BE((unsigned char *)header->par_utf16_name, &outlen, (const unsigned char *)par_filename,
                                (int *)&par_fn_len);
        if (utf_ret < 0) {
                mvhd_set_encoding_err(utf_ret, (int *)err);
                rv = -1;
                goto end;
        }

        /* And encode the paths to UTF16-LE */
        size_t par_path_len = strlen(par_path);
        outlen = sizeof *w2ku_path_buff * MVHD_MAX_PATH_CHARS;
        utf_ret = UTF8ToUTF16LE((unsigned char *)w2ku_path_buff, &outlen, (const unsigned char *)par_path, (int *)&par_path_len);
        if (utf_ret < 0) {
                mvhd_set_encoding_err(utf_ret, (int *)err);
                rv = -1;
                goto end;
        }
        int w2ku_len = utf_ret;
        outlen = sizeof *w2ru_path_buff * MVHD_MAX_PATH_CHARS;
        utf_ret = UTF8ToUTF16LE((unsigned char *)w2ru_path_buff, &outlen, (const unsigned char *)rel_path, (int *)&rel_len);
        if (utf_ret < 0) {
                mvhd_set_encoding_err(utf_ret, (int *)err);
                rv = -1;
                goto end;
        }
        int w2ru_len = utf_ret;
        /**
         * Finally populate the parent locaters in the sparse header.
         * This is the information needed to find the paths saved elsewhere
         * in the VHD image
         */

        /* Note about the plat_data_space field: The VHD spec says this field stores the number of sectors needed to store the
         * locator path. However, Hyper-V and VPC store the number of bytes, not the number of sectors, and will refuse to open
         * VHDs which have the number of sectors in this field. See
         * https://stackoverflow.com/questions/40760181/mistake-in-virtual-hard-disk-image-format-specification
         */
        header->par_loc_entry[0].plat_code = MVHD_DIF_LOC_W2KU;
        header->par_loc_entry[0].plat_data_len = (uint32_t)w2ku_len;
        header->par_loc_entry[0].plat_data_offset = (uint64_t)start_offset;
        header->par_loc_entry[0].plat_data_space =
                ((header->par_loc_entry[0].plat_data_len / MVHD_SECTOR_SIZE) + 1) * MVHD_SECTOR_SIZE;
        header->par_loc_entry[1].plat_code = MVHD_DIF_LOC_W2RU;
        header->par_loc_entry[1].plat_data_len = (uint32_t)w2ru_len;
        header->par_loc_entry[1].plat_data_offset = (uint64_t)start_offset + ((uint64_t)header->par_loc_entry[0].plat_data_space);
        header->par_loc_entry[1].plat_data_space =
                ((header->par_loc_entry[1].plat_data_len / MVHD_SECTOR_SIZE) + 1) * MVHD_SECTOR_SIZE;
        goto end;

end:
        return rv;
}

MVHDMeta *mvhd_create_fixed(const char *path, MVHDGeom geom, int *err, mvhd_progress_callback progress_callback) {
        uint64_t size_in_bytes = mvhd_calc_size_bytes(&geom);
        return mvhd_create_fixed_raw(path, NULL, size_in_bytes, &geom, err, progress_callback);
}

/**
 * \brief internal function that implements public mvhd_create_fixed() functionality
 *
 * Contains one more parameter than the public function, to allow using an existing
 * raw disk image as the data source for the new fixed VHD.
 *
 * \param [in] raw_image file handle to a raw disk image to populate VHD
 */
MVHDMeta *mvhd_create_fixed_raw(const char *path, FILE *raw_img, uint64_t size_in_bytes, MVHDGeom *geom, int *err,
                                mvhd_progress_callback progress_callback) {
        uint8_t img_data[MVHD_SECTOR_SIZE] = {0};
        uint8_t footer_buff[MVHD_FOOTER_SIZE] = {0};
        MVHDMeta *vhdm = calloc(1, sizeof *vhdm);
        if (vhdm == NULL) {
                *err = MVHD_ERR_MEM;
                goto end;
        }
        if (geom == NULL || (geom->cyl == 0 || geom->heads == 0 || geom->spt == 0)) {
                *err = MVHD_ERR_INVALID_GEOM;
                goto cleanup_vhdm;
        }
        FILE *f = mvhd_fopen(path, "wb+", err);
        if (f == NULL) {
                goto cleanup_vhdm;
        }
        mvhd_fseeko64(f, 0, SEEK_SET);
        uint32_t size_sectors = (uint32_t)(size_in_bytes / MVHD_SECTOR_SIZE);
        uint32_t s;
        if (progress_callback)
                progress_callback(0, size_sectors);
        if (raw_img != NULL) {
                mvhd_fseeko64(raw_img, 0, SEEK_END);
                uint64_t raw_size = (uint64_t)mvhd_ftello64(raw_img);
                MVHDGeom raw_geom = mvhd_calculate_geometry(raw_size);
                if (mvhd_calc_size_bytes(&raw_geom) != raw_size) {
                        *err = MVHD_ERR_CONV_SIZE;
                        goto cleanup_vhdm;
                }
                mvhd_gen_footer(&vhdm->footer, raw_size, geom, MVHD_TYPE_FIXED, 0);
                mvhd_fseeko64(raw_img, 0, SEEK_SET);
                for (s = 0; s < size_sectors; s++) {
                        fread(img_data, sizeof img_data, 1, raw_img);
                        fwrite(img_data, sizeof img_data, 1, f);
                        if (progress_callback)
                                progress_callback(s + 1, size_sectors);
                }
        } else {
                mvhd_gen_footer(&vhdm->footer, size_in_bytes, geom, MVHD_TYPE_FIXED, 0);
                for (s = 0; s < size_sectors; s++) {
                        fwrite(img_data, sizeof img_data, 1, f);
                        if (progress_callback)
                                progress_callback(s + 1, size_sectors);
                }
        }
        mvhd_footer_to_buffer(&vhdm->footer, footer_buff);
        fwrite(footer_buff, sizeof footer_buff, 1, f);
        fclose(f);
        f = NULL;
        free(vhdm);
        vhdm = mvhd_open(path, false, err);
        goto end;

cleanup_vhdm:
        free(vhdm);
        vhdm = NULL;
end:
        return vhdm;
}

/**
 * \brief Create sparse or differencing VHD image.
 *
 * \param [in] path is the absolute path to the VHD file to create
 * \param [in] par_path is the absolute path to a parent image. If NULL, a sparse image is created, otherwise create a
 * differencing image \param [in] size_in_bytes is the total size in bytes of the virtual hard disk image \param [in] geom is the
 * HDD geometry of the image to create. Determines final image size \param [in] block_size_in_sectors is the block size in sectors
 * \param [out] err indicates what error occurred, if any
 *
 * \return NULL if an error occurrs. Check value of *err for actual error. Otherwise returns pointer to a MVHDMeta struct
 */
static MVHDMeta *mvhd_create_sparse_diff(const char *path, const char *par_path, uint64_t size_in_bytes, MVHDGeom *geom,
                                         uint32_t block_size_in_sectors, int *err) {
        uint8_t footer_buff[MVHD_FOOTER_SIZE] = {0};
        uint8_t sparse_buff[MVHD_SPARSE_SIZE] = {0};
        uint8_t bat_sect[MVHD_SECTOR_SIZE];
        MVHDGeom par_geom = {0};
        memset(bat_sect, 0xffffffff, sizeof bat_sect);
        MVHDMeta *vhdm = NULL;
        MVHDMeta *par_vhdm = NULL;
        mvhd_utf16 *w2ku_path_buff = NULL;
        mvhd_utf16 *w2ru_path_buff = NULL;
        uint32_t par_mod_timestamp = 0;
        if (par_path != NULL) {
                par_mod_timestamp = mvhd_file_mod_timestamp(par_path, err);
                if (*err != 0) {
                        goto end;
                }
                par_vhdm = mvhd_open(par_path, true, err);
                if (par_vhdm == NULL) {
                        goto end;
                }
        }
        vhdm = calloc(1, sizeof *vhdm);
        if (vhdm == NULL) {
                *err = MVHD_ERR_MEM;
                goto cleanup_par_vhdm;
        }
        if (par_vhdm != NULL) {
                /* We use the geometry from the parent VHD, not what was passed in */
                par_geom.cyl = par_vhdm->footer.geom.cyl;
                par_geom.heads = par_vhdm->footer.geom.heads;
                par_geom.spt = par_vhdm->footer.geom.spt;
                geom = &par_geom;
                size_in_bytes = par_vhdm->footer.curr_sz;
        } else if (geom == NULL || (geom->cyl == 0 || geom->heads == 0 || geom->spt == 0)) {
                *err = MVHD_ERR_INVALID_GEOM;
                goto cleanup_vhdm;
        }

        FILE *f = mvhd_fopen(path, "wb+", err);
        if (f == NULL) {
                goto cleanup_vhdm;
        }
        mvhd_fseeko64(f, 0, SEEK_SET);
        /* Note, the sparse header follows the footer copy at the beginning of the file */
        if (par_path == NULL) {
                mvhd_gen_footer(&vhdm->footer, size_in_bytes, geom, MVHD_TYPE_DYNAMIC, MVHD_FOOTER_SIZE);
        } else {
                mvhd_gen_footer(&vhdm->footer, size_in_bytes, geom, MVHD_TYPE_DIFF, MVHD_FOOTER_SIZE);
        }
        mvhd_footer_to_buffer(&vhdm->footer, footer_buff);
        /* As mentioned, start with a copy of the footer */
        fwrite(footer_buff, sizeof footer_buff, 1, f);
        /**
         * Calculate the number of (2MB or 512KB) data blocks required to store the entire
         * contents of the disk image, followed by the number of sectors the
         * BAT occupies in the image. Note, the BAT is sector aligned, and is padded
         * to the next sector boundary
         * */
        uint32_t size_in_sectors = (uint32_t)(size_in_bytes / MVHD_SECTOR_SIZE);
        uint32_t num_blks = size_in_sectors / block_size_in_sectors;
        if (size_in_sectors % block_size_in_sectors != 0) {
                num_blks += 1;
        }
        uint32_t num_bat_sect = num_blks / MVHD_BAT_ENT_PER_SECT;
        if (num_blks % MVHD_BAT_ENT_PER_SECT != 0) {
                num_bat_sect += 1;
        }
        /* Storing the BAT directly following the footer and header */
        uint64_t bat_offset = MVHD_FOOTER_SIZE + MVHD_SPARSE_SIZE;
        uint64_t par_loc_offset = 0;

        /**
         * If creating a differencing VHD, populate the sparse header with additional
         * data about the parent image, and where to find it, and it's last modified timestamp
         * */
        if (par_vhdm != NULL) {
                /**
                 * Create output buffers to encode paths into.
                 * The paths are not stored directly in the sparse header, hence the need to
                 * store them in buffers to be written to the VHD image later
                 */
                w2ku_path_buff = calloc(MVHD_MAX_PATH_CHARS, sizeof *w2ku_path_buff);
                if (w2ku_path_buff == NULL) {
                        *err = MVHD_ERR_MEM;
                        goto end;
                }
                w2ru_path_buff = calloc(MVHD_MAX_PATH_CHARS, sizeof *w2ru_path_buff);
                if (w2ru_path_buff == NULL) {
                        *err = MVHD_ERR_MEM;
                        goto end;
                }
                memcpy(vhdm->sparse.par_uuid, par_vhdm->footer.uuid, sizeof vhdm->sparse.par_uuid);
                par_loc_offset = bat_offset + ((uint64_t)num_bat_sect * MVHD_SECTOR_SIZE) + (5 * MVHD_SECTOR_SIZE);
                if (mvhd_gen_par_loc(&vhdm->sparse, path, par_path, par_loc_offset, w2ku_path_buff, w2ru_path_buff,
                                     (MVHDError *)err) < 0) {
                        goto cleanup_vhdm;
                }
                vhdm->sparse.par_timestamp = par_mod_timestamp;
        }
        mvhd_gen_sparse_header(&vhdm->sparse, num_blks, bat_offset, block_size_in_sectors);
        mvhd_header_to_buffer(&vhdm->sparse, sparse_buff);
        fwrite(sparse_buff, sizeof sparse_buff, 1, f);
        /* The BAT sectors need to be filled with 0xffffffff */
        uint32_t i;
        for (i = 0; i < num_bat_sect; i++) {
                fwrite(bat_sect, sizeof bat_sect, 1, f);
        }
        mvhd_write_empty_sectors(f, 5);
        /**
         * If creating a differencing VHD, the paths to the parent image need to be written
         * tp the file. Both absolute and relative paths are written
         * */
        if (par_vhdm != NULL) {
                uint64_t curr_pos = (uint64_t)mvhd_ftello64(f);
                /* Double check my sums... */
                assert(curr_pos == par_loc_offset);
                /* Fill the space required for location data with zero */
                uint8_t empty_sect[MVHD_SECTOR_SIZE] = {0};
                int i;
                uint32_t j;
                for (i = 0; i < 2; i++) {
                        for (j = 0; j < (vhdm->sparse.par_loc_entry[i].plat_data_space / MVHD_SECTOR_SIZE); j++) {
                                fwrite(empty_sect, sizeof empty_sect, 1, f);
                        }
                }
                /* Now write the location entries */
                mvhd_fseeko64(f, vhdm->sparse.par_loc_entry[0].plat_data_offset, SEEK_SET);
                fwrite(w2ku_path_buff, vhdm->sparse.par_loc_entry[0].plat_data_len, 1, f);
                mvhd_fseeko64(f, vhdm->sparse.par_loc_entry[1].plat_data_offset, SEEK_SET);
                fwrite(w2ru_path_buff, vhdm->sparse.par_loc_entry[1].plat_data_len, 1, f);
                /* and reset the file position to continue */
                mvhd_fseeko64(f, vhdm->sparse.par_loc_entry[1].plat_data_offset + vhdm->sparse.par_loc_entry[1].plat_data_space,
                              SEEK_SET);
                mvhd_write_empty_sectors(f, 5);
        }
        /* And finish with the footer */
        fwrite(footer_buff, sizeof footer_buff, 1, f);
        fclose(f);
        f = NULL;
        free(vhdm);
        vhdm = mvhd_open(path, false, err);
        goto end;

cleanup_vhdm:
        free(vhdm);
        vhdm = NULL;
cleanup_par_vhdm:
        if (par_vhdm != NULL) {
                mvhd_close(par_vhdm);
        }
end:
        free(w2ku_path_buff);
        free(w2ru_path_buff);
        return vhdm;
}

MVHDMeta *mvhd_create_sparse(const char *path, MVHDGeom geom, int *err) {
        uint64_t size_in_bytes = mvhd_calc_size_bytes(&geom);
        return mvhd_create_sparse_diff(path, NULL, size_in_bytes, &geom, MVHD_BLOCK_LARGE, err);
}

MVHDMeta *mvhd_create_diff(const char *path, const char *par_path, int *err) {
        return mvhd_create_sparse_diff(path, par_path, 0, NULL, MVHD_BLOCK_LARGE, err);
}

MVHDMeta *mvhd_create_ex(MVHDCreationOptions options, int *err) {
        uint32_t geom_sector_size;
        switch (options.type) {
        case MVHD_TYPE_FIXED:
        case MVHD_TYPE_DYNAMIC:
                geom_sector_size = mvhd_calc_size_sectors(&(options.geometry));
                if ((options.size_in_bytes > 0 && (options.size_in_bytes % MVHD_SECTOR_SIZE) > 0) ||
                    (options.size_in_bytes > MVHD_MAX_SIZE_IN_BYTES) || (options.size_in_bytes == 0 && geom_sector_size == 0)) {
                        *err = MVHD_ERR_INVALID_SIZE;
                        return NULL;
                }

                if (options.size_in_bytes > 0 && ((uint64_t)geom_sector_size * MVHD_SECTOR_SIZE) > options.size_in_bytes) {
                        *err = MVHD_ERR_INVALID_GEOM;
                        return NULL;
                }

                if (options.size_in_bytes == 0)
                        options.size_in_bytes = (uint64_t)geom_sector_size * MVHD_SECTOR_SIZE;

                if (geom_sector_size == 0)
                        options.geometry = mvhd_calculate_geometry(options.size_in_bytes);
                break;
        case MVHD_TYPE_DIFF:
                if (options.parent_path == NULL) {
                        *err = MVHD_ERR_FILE;
                        return NULL;
                }
                break;
        default:
                *err = MVHD_ERR_TYPE;
                return NULL;
        }

        if (options.path == NULL) {
                *err = MVHD_ERR_FILE;
                return NULL;
        }

        if (options.type != MVHD_TYPE_FIXED) {
                if (options.block_size_in_sectors == MVHD_BLOCK_DEFAULT)
                        options.block_size_in_sectors = MVHD_BLOCK_LARGE;

                if (options.block_size_in_sectors != MVHD_BLOCK_LARGE && options.block_size_in_sectors != MVHD_BLOCK_SMALL) {
                        *err = MVHD_ERR_INVALID_BLOCK_SIZE;
                        return NULL;
                }
        }

        switch (options.type) {
        case MVHD_TYPE_FIXED:
                return mvhd_create_fixed_raw(options.path, NULL, options.size_in_bytes, &(options.geometry), err,
                                             options.progress_callback);
        case MVHD_TYPE_DYNAMIC:
                return mvhd_create_sparse_diff(options.path, NULL, options.size_in_bytes, &(options.geometry),
                                               options.block_size_in_sectors, err);
        case MVHD_TYPE_DIFF:
                return mvhd_create_sparse_diff(options.path, options.parent_path, 0, NULL, options.block_size_in_sectors, err);
        }

        return NULL; /* Make the compiler happy */
}
