#pragma once

#include <zstd.h>
#include "../data/mem.hpp"
#include "logger.hpp"

namespace zstd {

    /// Decompressor with ref to dictionnary
    class read_ctx {
    public:
        read_ctx(ZSTD_DDict *dict): ctx(ZSTD_createDCtx()), dict(dict) { }
        ~read_ctx() {
            ZSTD_freeDCtx(ctx);
        }

        /// Extract in to out (error message on failure)
        template<typename I, typename O>
        std::optional<const char*> decompress(const I& in, O& out) const {
            const auto maxSize = ZSTD_getFrameContentSize(in.data(), in.size());
            out.resize(maxSize);
            const auto actualSize = ZSTD_decompress_usingDDict(ctx, out.data(), out.size(), in.data(), in.size(), dict);
            if(ZSTD_isError(actualSize)) {
                return ZSTD_getErrorName(actualSize);
            }
            out.resize(actualSize);
            return std::nullopt;
        }

    protected:
        ZSTD_DCtx *ctx;
        ZSTD_DDict *dict;
    };

    /// Decompressor with builtin dictionnary
    class read_dict_ctx: public read_ctx {
    public:
        read_dict_ctx(const void* data, size_t size): read_ctx(ZSTD_createDDict(data, size)) {
            assert(dict != NULL);
        }
        ~read_dict_ctx() {
            ZSTD_freeDDict(dict);
        }
    };

    /// Compressor with ref to dictionnary
    class write_ctx {
    public:
        write_ctx(ZSTD_CDict *dict): ctx(ZSTD_createCCtx()), dict(dict) { }
        ~write_ctx() {
            ZSTD_freeCCtx(ctx);
        }

        /// Compress in to out (error message on failure)
        template<typename I, typename O>
        std::optional<const char*> compress(const I& in, O& out) const {
            const auto maxSize = ZSTD_compressBound(in.size());
            out.resize(maxSize);
            const auto actualSize = ZSTD_compress_usingCDict(ctx, out.data(), out.size(), in.data(), in.size(), dict);
            if(ZSTD_isError(actualSize)) {
                return ZSTD_getErrorName(actualSize);
            }
            out.resize(actualSize);
            return std::nullopt;
        }

    protected:
        ZSTD_CCtx *ctx;
        ZSTD_CDict *dict;
    };

    class dict_set {
    public:
        dict_set(const std::vector<char>& data) { load(data.data(), data.size()); }
        dict_set(const data::out_view& data) { load(data.data(), data.size()); }
        ~dict_set() {
            ZSTD_freeCDict(c);
            ZSTD_freeDDict(d);
        }

        read_ctx make_reader() const {
            return read_ctx(d);
        }
        write_ctx make_writer() const {
            return write_ctx(c);
        }
    private:
        void load(const void* data, size_t size) {
            c = ZSTD_createCDict(data, size, ZSTD_CLEVEL_DEFAULT);
            assert(c != NULL);
            d = ZSTD_createDDict(data, size);
            assert(d != NULL);
        }

        ZSTD_CDict *c;
        ZSTD_DDict *d;
    };
}