#include "tar_archive.h"

#include <microtar.h>
#include <storage/storage.h>
#include <furi.h>
#include <toolbox/path.h>

#include <lib/uzlib/src/uzlib.h>

typedef struct {
    File* file;
    uint8_t* buffer;
    size_t buffer_size;
    uint8_t* dict;
    size_t dict_size;
    struct uzlib_uncomp uzlib;

    uint32_t source_pos;
    uint32_t dest_pos;
    bool eof;
} Gunzip;

int gunzip_read_cb(struct uzlib_uncomp* uncomp) {
    void* data_p = uncomp;
    data_p -= offsetof(Gunzip, uzlib);
    Gunzip* gunzip = data_p;

    uint16_t read_size = storage_file_read(gunzip->file, gunzip->buffer, gunzip->buffer_size);
    gunzip->uzlib.source = &gunzip->buffer[1]; // we will return buffer[0] at exit
    gunzip->uzlib.source_limit = gunzip->buffer + read_size;

    if(read_size == 0) {
        return -1;
    }
    gunzip->source_pos += read_size;

    return gunzip->buffer[0];
}

Gunzip* gunzip_alloc(File* file, size_t dict_size, size_t buffer_size) {
    Gunzip* gunzip = malloc(sizeof(Gunzip));
    gunzip->file = file;
    gunzip->buffer_size = buffer_size;
    gunzip->buffer = malloc(buffer_size);
    gunzip->dict_size = dict_size;
    gunzip->dict = malloc(dict_size);

    uzlib_uncompress_init(&gunzip->uzlib, gunzip->dict, gunzip->dict_size);

    gunzip->uzlib.source = 0;
    gunzip->uzlib.source_limit = 0;
    gunzip->uzlib.source_read_cb = gunzip_read_cb;
    gunzip->source_pos = 0;
    gunzip->dest_pos = 0;
    gunzip->eof = false;

    return gunzip;
}

void gunzip_free(Gunzip* gunzip) {
    free(gunzip->buffer);
    free(gunzip->dict);
    free(gunzip);
}

int32_t gunzip_uncompress(Gunzip* gunzip, void* out, size_t out_len) {
    if(gunzip->eof) {
        return 0;
    }

    gunzip->uzlib.dest = out;
    gunzip->uzlib.dest_limit = (uint8_t*)out + out_len;

    int res = uzlib_uncompress_chksum(&gunzip->uzlib);

    if(res == TINF_DONE) {
        gunzip->eof = true;
    }
    if(res < 0) {
        return res;
    }
    int32_t read = gunzip->uzlib.dest - (uint8_t*)out;
    gunzip->dest_pos += read;
    return read;
}

int32_t gunzip_seek(Gunzip* gunzip, size_t pos) {
    if(pos == gunzip->dest_pos) {
        return 0;
    }

    if(pos < gunzip->dest_pos) {
        // TODO: rewind to start if this causes issues
        furi_crash("Gunzip rewind");
        return -1;
    }

    size_t void_size = MIN(4096U, pos - gunzip->dest_pos);
    void* void_buf = malloc(void_size);
    while(!gunzip->eof && gunzip->dest_pos < pos) {
        size_t uncompress_size = MIN(void_size, pos - gunzip->dest_pos);
        int res = gunzip_uncompress(gunzip, void_buf, uncompress_size);
        if(res < 0) {
            free(void_buf);
            return res;
        }
    }
    free(void_buf);
    return (pos == gunzip->dest_pos) ? 0 : -1;
}

#define TAG "TarArch"
#define MAX_NAME_LEN 254
#define FILE_BLOCK_SIZE 512

#define FILE_OPEN_NTRIES 10
#define FILE_OPEN_RETRY_DELAY 25

typedef struct TarArchive {
    Storage* storage;
    mtar_t tar;
    tar_unpack_file_cb unpack_cb;
    void* unpack_cb_context;

    tar_unpack_read_cb read_cb;
    void* read_cb_context;
    size_t total_size;
    Gunzip* gunzip;
} TarArchive;

/* API WRAPPER */
static int mtar_storage_file_write(void* stream, const void* data, unsigned size) {
    uint16_t bytes_written = storage_file_write(stream, data, size);
    return (bytes_written == size) ? bytes_written : MTAR_EWRITEFAIL;
}

static int mtar_storage_file_read(void* stream, void* data, unsigned size) {
    uint16_t bytes_read = storage_file_read(stream, data, size);
    return (bytes_read == size) ? bytes_read : MTAR_EREADFAIL;
}

static int mtar_storage_file_seek(void* stream, unsigned offset) {
    bool res = storage_file_seek(stream, offset, true);
    return res ? MTAR_ESUCCESS : MTAR_ESEEKFAIL;
}

static int mtar_storage_file_close(void* stream) {
    if(stream) {
        storage_file_close(stream);
        storage_file_free(stream);
    }
    return MTAR_ESUCCESS;
}

const struct mtar_ops filesystem_ops = {
    .read = mtar_storage_file_read,
    .write = mtar_storage_file_write,
    .seek = mtar_storage_file_seek,
    .close = mtar_storage_file_close,
};

static int mtar_storage_gunzip_write(void* gunzip, const void* data, unsigned size) {
    UNUSED(gunzip);
    UNUSED(data);
    UNUSED(size);
    furi_crash("Write to gzipped tar");
    return MTAR_EWRITEFAIL;
}

static int mtar_storage_gunzip_read(void* gunzip, void* data, unsigned size) {
    int32_t res = gunzip_uncompress(gunzip, data, size);
    if(res < 0) {
        FURI_LOG_E(TAG, "Error uncompressing gzip: %ld\n", res);
    }
    return (res == (int32_t)size) ? res : MTAR_EREADFAIL;
}

static int mtar_storage_gunzip_seek(void* gunzip, unsigned offset) {
    int32_t res = gunzip_seek(gunzip, offset);
    if(res < 0) {
        FURI_LOG_E(TAG, "Error seeking gzip: %ld\n", res);
    }
    return (res == 0) ? MTAR_ESUCCESS : MTAR_ESEEKFAIL;
}

static int mtar_storage_gunzip_close(void* _gunzip) {
    Gunzip* gunzip = _gunzip;
    if(gunzip) {
        if(gunzip->file) {
            storage_file_close(gunzip->file);
            storage_file_free(gunzip->file);
        }
        gunzip_free(gunzip);
    }
    return MTAR_ESUCCESS;
}

const struct mtar_ops gunzip_ops = {
    .read = mtar_storage_gunzip_read,
    .write = mtar_storage_gunzip_write,
    .seek = mtar_storage_gunzip_seek,
    .close = mtar_storage_gunzip_close,
};

TarArchive* tar_archive_alloc(Storage* storage) {
    furi_check(storage);
    TarArchive* archive = malloc(sizeof(TarArchive));
    archive->storage = storage;
    archive->unpack_cb = NULL;
    return archive;
}

bool tar_archive_open(TarArchive* archive, const char* path, TarOpenMode mode) {
    furi_check(archive);
    FS_AccessMode access_mode;
    FS_OpenMode open_mode;
    int mtar_access = 0;

    switch(mode) {
    case TAR_OPEN_MODE_READ:
        mtar_access = MTAR_READ;
        access_mode = FSAM_READ;
        open_mode = FSOM_OPEN_EXISTING;
        break;
    case TAR_OPEN_MODE_WRITE:
        mtar_access = MTAR_WRITE;
        access_mode = FSAM_WRITE;
        open_mode = FSOM_CREATE_ALWAYS;
        break;
    default:
        return false;
    }

    File* stream = storage_file_alloc(archive->storage);
    if(!storage_file_open(stream, path, access_mode, open_mode)) {
        storage_file_free(stream);
        return false;
    }
    archive->total_size = storage_file_size(stream);

    char* dot = strrchr(path, '.');
    if(dot == NULL || strcmp(dot, ".gz") != 0 || mode != TAR_OPEN_MODE_READ) {
        mtar_init(&archive->tar, mtar_access, &filesystem_ops, stream);
    } else {
        archive->gunzip = gunzip_alloc(stream, 32 * 1024, 10 * 1024);

        int res = uzlib_gzip_parse_header(&archive->gunzip->uzlib);
        if(res != TINF_OK) {
            FURI_LOG_E(TAG, "Error parsing gzip header: %d\n", res);
            storage_file_close(stream);
            storage_file_free(stream);
            gunzip_free(archive->gunzip);
            archive->gunzip = NULL;
            return false;
        }

        mtar_init(&archive->tar, mtar_access, &gunzip_ops, archive->gunzip);
    }

    return true;
}

void tar_archive_free(TarArchive* archive) {
    furi_check(archive);
    if(mtar_is_open(&archive->tar)) {
        mtar_close(&archive->tar);
    }
    free(archive);
}

void tar_archive_set_file_callback(TarArchive* archive, tar_unpack_file_cb callback, void* context) {
    furi_check(archive);
    archive->unpack_cb = callback;
    archive->unpack_cb_context = context;
}

void tar_archive_set_read_callback(TarArchive* archive, tar_unpack_read_cb callback, void* context) {
    furi_check(archive);
    archive->read_cb = callback;
    archive->read_cb_context = context;
}

static int tar_archive_entry_counter(mtar_t* tar, const mtar_header_t* header, void* param) {
    UNUSED(tar);
    UNUSED(header);
    furi_assert(param);
    int32_t* counter = param;
    (*counter)++;
    return 0;
}

int32_t tar_archive_get_entries_count(TarArchive* archive) {
    furi_check(archive);
    int32_t counter = 0;
    if(mtar_foreach(&archive->tar, tar_archive_entry_counter, &counter) != MTAR_ESUCCESS) {
        counter = -1;
    }
    return counter;
}

bool tar_archive_dir_add_element(TarArchive* archive, const char* dirpath) {
    furi_check(archive);
    return (mtar_write_dir_header(&archive->tar, dirpath) == MTAR_ESUCCESS);
}

bool tar_archive_finalize(TarArchive* archive) {
    furi_check(archive);
    return (mtar_finalize(&archive->tar) == MTAR_ESUCCESS);
}

bool tar_archive_store_data(
    TarArchive* archive,
    const char* path,
    const uint8_t* data,
    const int32_t data_len) {
    furi_check(archive);

    return (
        tar_archive_file_add_header(archive, path, data_len) &&
        tar_archive_file_add_data_block(archive, data, data_len) &&
        tar_archive_file_finalize(archive));
}

bool tar_archive_file_add_header(TarArchive* archive, const char* path, const int32_t data_len) {
    furi_check(archive);

    return (mtar_write_file_header(&archive->tar, path, data_len) == MTAR_ESUCCESS);
}

bool tar_archive_file_add_data_block(
    TarArchive* archive,
    const uint8_t* data_block,
    const int32_t block_len) {
    furi_check(archive);

    return (mtar_write_data(&archive->tar, data_block, block_len) == block_len);
}

bool tar_archive_file_finalize(TarArchive* archive) {
    furi_check(archive);
    return (mtar_end_data(&archive->tar) == MTAR_ESUCCESS);
}

typedef struct {
    TarArchive* archive;
    const char* work_dir;
    Storage_name_converter converter;
} TarArchiveDirectoryOpParams;

static bool archive_extract_current_file(TarArchive* archive, const char* dst_path) {
    mtar_t* tar = &archive->tar;
    File* out_file = storage_file_alloc(archive->storage);
    uint8_t* readbuf = malloc(FILE_BLOCK_SIZE);

    bool success = true;
    uint8_t n_tries = FILE_OPEN_NTRIES;
    do {
        while(n_tries-- > 0) {
            if(storage_file_open(out_file, dst_path, FSAM_WRITE, FSOM_CREATE_ALWAYS)) {
                break;
            }
            FURI_LOG_W(TAG, "Failed to open '%s', reties: %d", dst_path, n_tries);
            storage_file_close(out_file);
            furi_delay_ms(FILE_OPEN_RETRY_DELAY);
        }

        if(!storage_file_is_open(out_file)) {
            success = false;
            break;
        }

        while(!mtar_eof_data(tar)) {
            int32_t readcnt = mtar_read_data(tar, readbuf, FILE_BLOCK_SIZE);
            if(!readcnt || !storage_file_write(out_file, readbuf, readcnt)) {
                success = false;
                break;
            }

            if(archive->read_cb) {
                archive->read_cb(
                    archive->gunzip ? archive->gunzip->source_pos : archive->tar.pos,
                    archive->total_size,
                    archive->read_cb_context);
            }
        }
    } while(false);
    storage_file_free(out_file);
    free(readbuf);

    return success;
}

static int archive_extract_foreach_cb(mtar_t* tar, const mtar_header_t* header, void* param) {
    UNUSED(tar);
    TarArchiveDirectoryOpParams* op_params = param;
    TarArchive* archive = op_params->archive;

    bool skip_entry = false;
    if(archive->unpack_cb) {
        skip_entry = !archive->unpack_cb(
            header->name, header->type == MTAR_TDIR, archive->unpack_cb_context);
    }

    if(skip_entry) {
        FURI_LOG_W(TAG, "filter: skipping entry \"%s\"", header->name);
        return 0;
    }

    FuriString* full_extracted_fname;
    if(header->type == MTAR_TDIR) {
        full_extracted_fname = furi_string_alloc();
        path_concat(op_params->work_dir, header->name, full_extracted_fname);

        bool create_res =
            storage_simply_mkdir(archive->storage, furi_string_get_cstr(full_extracted_fname));
        furi_string_free(full_extracted_fname);
        return create_res ? 0 : -1;
    }

    if(header->type != MTAR_TREG) {
        FURI_LOG_W(TAG, "not extracting unsupported type \"%s\"", header->name);
        return 0;
    }

    FURI_LOG_D(TAG, "Extracting %u bytes to '%s'", header->size, header->name);

    FuriString* converted_fname = furi_string_alloc_set(header->name);
    if(op_params->converter) {
        op_params->converter(converted_fname);
    }

    full_extracted_fname = furi_string_alloc();
    path_concat(op_params->work_dir, furi_string_get_cstr(converted_fname), full_extracted_fname);

    bool success =
        archive_extract_current_file(archive, furi_string_get_cstr(full_extracted_fname));

    furi_string_free(converted_fname);
    furi_string_free(full_extracted_fname);
    return success ? 0 : -1;
}

bool tar_archive_unpack_to(
    TarArchive* archive,
    const char* destination,
    Storage_name_converter converter) {
    furi_check(archive);
    TarArchiveDirectoryOpParams param = {
        .archive = archive,
        .work_dir = destination,
        .converter = converter,
    };

    FURI_LOG_I(TAG, "Restoring '%s'", destination);

    return (mtar_foreach(&archive->tar, archive_extract_foreach_cb, &param) == MTAR_ESUCCESS);
};

bool tar_archive_add_file(
    TarArchive* archive,
    const char* fs_file_path,
    const char* archive_fname,
    const int32_t file_size) {
    furi_check(archive);
    uint8_t* file_buffer = malloc(FILE_BLOCK_SIZE);
    bool success = false;
    File* src_file = storage_file_alloc(archive->storage);
    uint8_t n_tries = FILE_OPEN_NTRIES;
    do {
        while(n_tries-- > 0) {
            if(storage_file_open(src_file, fs_file_path, FSAM_READ, FSOM_OPEN_EXISTING)) {
                break;
            }
            FURI_LOG_W(TAG, "Failed to open '%s', reties: %d", fs_file_path, n_tries);
            storage_file_close(src_file);
            furi_delay_ms(FILE_OPEN_RETRY_DELAY);
        }

        if(!storage_file_is_open(src_file) ||
           !tar_archive_file_add_header(archive, archive_fname, file_size)) {
            break;
        }

        success = true; // if file is empty, that's not an error
        uint16_t bytes_read = 0;
        while((bytes_read = storage_file_read(src_file, file_buffer, FILE_BLOCK_SIZE))) {
            success = tar_archive_file_add_data_block(archive, file_buffer, bytes_read);
            if(!success) {
                break;
            }
        }

        success = success && tar_archive_file_finalize(archive);
    } while(false);

    storage_file_free(src_file);
    free(file_buffer);
    return success;
}

bool tar_archive_add_dir(TarArchive* archive, const char* fs_full_path, const char* path_prefix) {
    furi_check(archive);
    furi_check(path_prefix);

    File* directory = storage_file_alloc(archive->storage);
    FileInfo file_info;

    FURI_LOG_I(TAG, "Backing up '%s', '%s'", fs_full_path, path_prefix);
    char* name = malloc(MAX_NAME_LEN);
    bool success = false;

    do {
        if(!storage_dir_open(directory, fs_full_path)) {
            break;
        }

        while(true) {
            if(!storage_dir_read(directory, &file_info, name, MAX_NAME_LEN)) {
                success = true; /* empty dir / no more files */
                break;
            }

            FuriString* element_name = furi_string_alloc();
            FuriString* element_fs_abs_path = furi_string_alloc();

            path_concat(fs_full_path, name, element_fs_abs_path);
            if(strlen(path_prefix)) {
                path_concat(path_prefix, name, element_name);
            } else {
                furi_string_set(element_name, name);
            }

            if(file_info_is_dir(&file_info)) {
                success =
                    tar_archive_dir_add_element(archive, furi_string_get_cstr(element_name)) &&
                    tar_archive_add_dir(
                        archive,
                        furi_string_get_cstr(element_fs_abs_path),
                        furi_string_get_cstr(element_name));
            } else {
                success = tar_archive_add_file(
                    archive,
                    furi_string_get_cstr(element_fs_abs_path),
                    furi_string_get_cstr(element_name),
                    file_info.size);
            }
            furi_string_free(element_name);
            furi_string_free(element_fs_abs_path);

            if(!success) {
                break;
            }
        }
    } while(false);

    free(name);
    storage_file_free(directory);
    return success;
}

bool tar_archive_unpack_file(
    TarArchive* archive,
    const char* archive_fname,
    const char* destination) {
    furi_check(archive);
    furi_check(archive_fname);
    furi_check(destination);
    if(mtar_find(&archive->tar, archive_fname) != MTAR_ESUCCESS) {
        return false;
    }
    return archive_extract_current_file(archive, destination);
}