// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "util/zlib.h"

#include <butil/macros.h>

#include <zlib.h>

#include <cstring>

#include <memory>

#include <ostream>

#include <string>

#include "common/cast_set.h"

#define ZRETURN_NOT_OK(call) RETURN_IF_ERROR(ZlibResultToStatus(call))

namespace doris {

using namespace ErrorCode;

namespace zlib {

#include "common/compile_check_begin.h"

namespace {

Status ZlibResultToStatus(int rc) {

    switch (rc) {

    case Z_OK:

        return Status::OK();

    case Z_STREAM_END:

        return Status::EndOfFile("zlib EOF");

    case Z_NEED_DICT:

        return Status::Corruption("zlib error: NEED_DICT");

    case Z_ERRNO:

        return Status::IOError("zlib error: Z_ERRNO");

    case Z_STREAM_ERROR:

        return Status::Corruption("zlib error: STREAM_ERROR");

    case Z_DATA_ERROR:

        return Status::Corruption("zlib error: DATA_ERROR");

    case Z_MEM_ERROR:

        return Status::RuntimeError("zlib error: MEM_ERROR");

    case Z_BUF_ERROR:

        return Status::RuntimeError("zlib error: BUF_ERROR");

    case Z_VERSION_ERROR:

        return Status::RuntimeError("zlib error: VERSION_ERROR");

    default:

        return Status::RuntimeError("zlib error: unknown error {}", rc);

    }

}

} // anonymous namespace

Status Compress(Slice input, std::ostream* out) {

    return CompressLevel(input, Z_DEFAULT_COMPRESSION, out);

}

Status CompressLevel(Slice input, int level, std::ostream* out) {

    z_stream zs;

    memset(&zs, 0, sizeof(zs));

    ZRETURN_NOT_OK(deflateInit2(&zs, level, Z_DEFLATED, 15 + 16 /* 15 window bits, enable gzip */,

                                8 /* memory level, max is 9 */, Z_DEFAULT_STRATEGY));

    zs.avail_in = cast_set<unsigned int>(input.get_size());

    zs.next_in = (unsigned char*)(input.mutable_data());

    const int kChunkSize = 256 * 1024;

    std::unique_ptr<unsigned char[]> chunk(new unsigned char[kChunkSize]);

    int flush;

    do {

        zs.avail_out = kChunkSize;

        zs.next_out = chunk.get();

        flush = (zs.avail_in == 0) ? Z_FINISH : Z_NO_FLUSH;

        Status s = ZlibResultToStatus(deflate(&zs, flush));

        if (!s.ok() && !s.is<END_OF_FILE>()) {

            return s;

        }

        int out_size = cast_set<int>(zs.next_out - chunk.get());

        if (out_size > 0) {

            out->write(reinterpret_cast<char*>(chunk.get()), out_size);

        }

    } while (flush != Z_FINISH);

    ZRETURN_NOT_OK(deflateEnd(&zs));

    return Status::OK();

}

Status Uncompress(Slice compressed, std::ostream* out) {

    z_stream zs;

    memset(&zs, 0, sizeof(zs));

    zs.next_in = (unsigned char*)(compressed.mutable_data());

    zs.avail_in = cast_set<unsigned int>(compressed.get_size());

    ZRETURN_NOT_OK(inflateInit2(&zs, 15 + 16 /* 15 window bits, enable zlib */));

    int flush;

    Status s;

    do {

        unsigned char buf[4096];

        zs.next_out = buf;

        zs.avail_out = arraysize(buf);

        flush = zs.avail_in > 0 ? Z_NO_FLUSH : Z_FINISH;

        s = ZlibResultToStatus(inflate(&zs, flush));

        if (!s.ok() && !s.is<END_OF_FILE>()) {

            return s;

        }

        out->write(reinterpret_cast<char*>(buf), zs.next_out - buf);

    } while (flush == Z_NO_FLUSH);

    ZRETURN_NOT_OK(inflateEnd(&zs));

    return Status::OK();

}

#include "common/compile_check_end.h"

} // namespace zlib

} // namespace doris