#include <string>
#include <vector>
#include "ppt.h"
#include "miniz.h"

//Some useful reads: 
//https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT
//https://weber.itn.liu.se/~stegu/OOXML/ISO_IEC_DIS_29500/C045515e/C045515e_PDF/Office_Open_XML_Part2_OpenPackagingConventions.pdf
//Note: For some reason NO_COMPRESSION setting does not produce a ZIP file that PPT can read.
//So for now we also compress PNG files.


std::string OpenZip(ZipArchive archive, mz_zip_archive &zip, std::string action) {
    if (archive.is_file() &&  archive.fname().empty()) return "Empty file name when "+action+" files.";
    if (!archive.is_file() &&  archive.mem().empty()) return "Empty archive when "+action+" files.";
    MZ_CLEAR_OBJ(zip);
    if (archive.is_file() 
            ? !mz_zip_reader_init_file(&zip, archive.fname().c_str(), MZ_ZIP_FLAG_DO_NOT_SORT_CENTRAL_DIRECTORY) 
            : !mz_zip_reader_init_mem(&zip, archive.mem().data(), archive.mem().size(), MZ_ZIP_FLAG_DO_NOT_SORT_CENTRAL_DIRECTORY))
        return "Could not read from "+archive.qfname_or_clipboard()+".";
    return {};
}


std::string GetFilesInZip(ZipArchive archive, std::vector<FileInZip>& files) {
    files.clear();
    mz_zip_archive zip_archive;
    std::string err = OpenZip(archive, zip_archive, "listing");
    if (err.size()) return err;
    files.reserve(mz_zip_reader_get_num_files(&zip_archive));
    for (unsigned i = 0; i<mz_zip_reader_get_num_files(&zip_archive); i++)
        if (mz_zip_archive_file_stat stat; !mz_zip_reader_file_stat(&zip_archive, i, &stat)) {
            mz_zip_reader_end(&zip_archive);
            return "Could not stat file #"+std::to_string(i) + " in " + archive.qfname_or_clipboard()+".";
        } else
            files.emplace_back(stat.m_filename, "");
    mz_zip_reader_end(&zip_archive);
    _ASSERT(err.empty());
    return err;
}

std::string ExtractFromZip(ZipArchive archive, std::vector<FileInZip>& files) {
    if (files.empty()) return "";
    mz_zip_archive zip_archive;
    std::string err = OpenZip(archive, zip_archive, "extracting");
    if (err.size()) return err;
    std::vector<FileInZip*> remove;
    for (FileInZip& f : files) {
        const int file_index = mz_zip_reader_locate_file(&zip_archive, f.inner_filename.c_str(), NULL, 0);
        if (file_index>=0) {
            size_t size = 0;
            void *p = mz_zip_reader_extract_to_heap(&zip_archive, file_index, &size, 0);
            if (p==nullptr) {
                mz_zip_reader_end(&zip_archive);
                return "Could not extract '"+f.inner_filename+"' from "+archive.qfname_or_clipboard()+".";
            } else {
                f.content.assign((char*)p, size);
                zip_archive.m_pFree(zip_archive.m_pAlloc_opaque, p);
            }
        } else
            remove.push_back(&f);
    }
    std::erase_if(files, [&remove](FileInZip& f) { return std::ranges::find(remove, &f)!=remove.end(); });
    mz_zip_reader_end(&zip_archive);
    _ASSERT(err.empty());
    return err;
}

size_t string_writer(void* pOpaque, mz_uint64 file_ofs, const void* pBuf, size_t n) {
    std::string& mem = *(std::string*)pOpaque;
    mem.resize(size_t(file_ofs));
    mem.append(static_cast<const char*>(pBuf), n); 
    return n;
}

std::string AddToZip(ZipArchive archive, const std::vector<FileInZip>& files_) {
    if (files_.empty()) return "";
    mz_zip_archive zip_archive;
    std::string err = OpenZip(archive, zip_archive, "updating");
    if (err.size()) return err;
    mz_zip_archive tmp_zip_archive, *use_for_appending = &zip_archive;
    std::string tmp_filename, tmp_content;
    std::vector<const FileInZip*> files_to_do;
    files_to_do.reserve(files_.size());
    for (const FileInZip& f: files_)
        files_to_do.push_back(&f);
    //If any of the files is already present, we need to copy the archive to remove them
    //Likewise if we work in memory: miniz cannot grow memory it has not allocated //TODO: fiddle with user re-alloc
    if (!archive.is_file() 
        || std::ranges::any_of(files_to_do, [&zip_archive](const FileInZip* f) { return mz_zip_reader_locate_file(&zip_archive, f->inner_filename.c_str(), NULL, 0)>=0; })) {
        MZ_CLEAR_OBJ(tmp_zip_archive);
        if (archive.is_file()) {
            tmp_filename = my_tmpnam();
            if (!mz_zip_writer_init_file(&tmp_zip_archive, tmp_filename.c_str(), 0)) {
                mz_zip_reader_end(&zip_archive);
                return StrCat("Could not open temporary file '", tmp_filename, "' for writing.");
            }
        } else {
            tmp_zip_archive.m_pWrite = &string_writer;
            tmp_zip_archive.m_pIO_opaque = &tmp_content;
            if (!mz_zip_writer_init(&tmp_zip_archive, 0)) {
                mz_zip_reader_end(&zip_archive);
                return "Could not open temporary in-mem archive.";
            }
        }
        std::string err_msg;
        for (unsigned i = 0; i<mz_zip_reader_get_num_files(&zip_archive); i++)
            if (mz_zip_archive_file_stat stat; !mz_zip_reader_file_stat(&zip_archive, i, &stat)) {
                err_msg = "Could not stat file #"+std::to_string(i)+" in "+archive.qfname_or_clipboard()+".";
            fail1:
                mz_zip_writer_end(&tmp_zip_archive);
                mz_zip_reader_end(&zip_archive);
                std::error_code ec;
                std::filesystem::remove(tmp_filename, ec);
                return err_msg;
            } else if (auto f = std::ranges::find_if(files_to_do, [&stat](const FileInZip* f) { return f && stat.m_filename==f->inner_filename; }); f!=files_to_do.end()) {
                //Replace this file to the one we have
                if (!mz_zip_writer_add_mem(&tmp_zip_archive, (*f)->inner_filename.c_str(), (*f)->content.data(), (*f)->content.size(), MZ_BEST_SPEED)) {
                    err_msg = "Cannot append file '" + (*f)->inner_filename + "' to archive when appending to "+archive.qfname_or_clipboard()+".";
                    goto fail1;
                }
                *f = nullptr; //dont append at the end
            } else if (!mz_zip_writer_add_from_zip_reader(&tmp_zip_archive, &zip_archive, i)) {
                err_msg = "Could not copy ZIP content to temporary " + (archive.is_file() ? "file '"+tmp_filename+"'" : "memory area") + " from "+archive.qfname_or_clipboard()+".";
                goto fail1;
            }
        use_for_appending = &tmp_zip_archive;
        mz_zip_reader_end(&zip_archive);
    } else
        mz_zip_writer_init_from_reader(&zip_archive, archive.is_file() ? archive.fname().c_str() : nullptr); //for files we need to supply the name to reopen writing
    for (auto* f : files_to_do)
        if (f && !mz_zip_writer_add_mem(use_for_appending, f->inner_filename.c_str(), f->content.data(), f->content.size(), MZ_BEST_SPEED)) {
            mz_zip_writer_end(use_for_appending);
            std::error_code ec;
            std::filesystem::remove(tmp_filename, ec);
            return "Cannot append file '" + f->inner_filename + "' to archive when appending to "+archive.qfname_or_clipboard()+".";
        }
    mz_zip_writer_finalize_archive(use_for_appending);
    mz_zip_writer_end(use_for_appending);
    //If we were appending in-place, we are done
    if (use_for_appending==&zip_archive) return {};
    if (archive.is_file()) {
        //if we used a temp file, copy the result over the original
        std::error_code ec;
        std::filesystem::rename(tmp_filename, archive.fname(), ec);
        if (ec) {
            std::filesystem::copy(tmp_filename, archive.fname(), ec);
            std::error_code ec2;
            std::filesystem::remove(tmp_filename, ec2);
            if (ec)
                return StrCat("Could not copy '", tmp_filename, "' to '", archive.fname(), "' after replacing files.");
        }
    } else {
        archive.mem().swap(tmp_content);
    }
    return {};
}
