/******************************************************************************
 *
 * Project:  GDAL
 * Purpose:  Zarr driver, ZarrV3CodecSequence class
 * Author:   Even Rouault <even dot rouault at spatialys.com>
 *
 ******************************************************************************
 * Copyright (c) 2023, Even Rouault <even dot rouault at spatialys.com>
 *
 * SPDX-License-Identifier: MIT
 ****************************************************************************/

#include "zarr_v3_codec.h"

/************************************************************************/
/*                     ZarrV3CodecSequence::Clone()                     */
/************************************************************************/

std::unique_ptr<ZarrV3CodecSequence> ZarrV3CodecSequence::Clone() const
{
    auto poClone = std::make_unique<ZarrV3CodecSequence>(m_oInputArrayMetadata);
    for (const auto &poCodec : m_apoCodecs)
        poClone->m_apoCodecs.emplace_back(poCodec->Clone());
    poClone->m_oCodecArray = m_oCodecArray.Clone();
    poClone->m_bPartialDecodingPossible = m_bPartialDecodingPossible;
    return poClone;
}

/************************************************************************/
/*                 ZarrV3CodecSequence::InitFromJson()                  */
/************************************************************************/

bool ZarrV3CodecSequence::InitFromJson(const CPLJSONObject &oCodecs,
                                       ZarrArrayMetadata &oOutputArrayMetadata)
{
    if (oCodecs.GetType() != CPLJSONObject::Type::Array)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "codecs is not an array");
        return false;
    }
    auto oCodecsArray = oCodecs.ToArray();

    ZarrArrayMetadata oInputArrayMetadata = m_oInputArrayMetadata;
    ZarrV3Codec::IOType eLastType = ZarrV3Codec::IOType::ARRAY;
    std::string osLastCodec;

    const auto InsertImplicitEndianCodecIfNeeded =
        [this, &oInputArrayMetadata, &eLastType, &osLastCodec]()
    {
        CPL_IGNORE_RET_VAL(this);
        if (eLastType == ZarrV3Codec::IOType::ARRAY &&
            oInputArrayMetadata.oElt.nativeSize > 1 &&
            oInputArrayMetadata.oElt.nativeType !=
                DtypeElt::NativeType::STRING_ASCII &&
            oInputArrayMetadata.oElt.nativeType !=
                DtypeElt::NativeType::STRING_UNICODE)
        {
            CPLError(CE_Warning, CPLE_AppDefined,
                     "'bytes' codec missing. Assuming little-endian storage, "
                     "but such tolerance may be removed in future versions");
            auto poEndianCodec = std::make_unique<ZarrV3CodecBytes>();
            ZarrArrayMetadata oTmpOutputArrayMetadata;
            poEndianCodec->InitFromConfiguration(
                ZarrV3CodecBytes::GetConfiguration(true), oInputArrayMetadata,
                oTmpOutputArrayMetadata, /* bEmitWarnings = */ true);
            oInputArrayMetadata = std::move(oTmpOutputArrayMetadata);
            eLastType = poEndianCodec->GetOutputType();
            osLastCodec = poEndianCodec->GetName();
            if constexpr (!CPL_IS_LSB)
            {
                // Insert a little endian codec if we are on a big endian target
                m_apoCodecs.emplace_back(std::move(poEndianCodec));
            }
        }
    };

    bool bShardingFound = false;
    std::vector<size_t> anBlockSizesBeforeSharding;
    for (const auto &oCodec : oCodecsArray)
    {
        if (oCodec.GetType() != CPLJSONObject::Type::Object)
        {
            CPLError(CE_Failure, CPLE_AppDefined, "codecs[] is not an object");
            return false;
        }
        const auto osName = oCodec["name"].ToString();
        std::unique_ptr<ZarrV3Codec> poCodec;
        if (osName == ZarrV3CodecGZip::NAME)
            poCodec = std::make_unique<ZarrV3CodecGZip>();
        else if (osName == ZarrV3CodecBlosc::NAME)
            poCodec = std::make_unique<ZarrV3CodecBlosc>();
        else if (osName == ZarrV3CodecZstd::NAME)
            poCodec = std::make_unique<ZarrV3CodecZstd>();
        else if (osName == ZarrV3CodecBytes::NAME ||
                 osName == "endian" /* endian is the old name */)
            poCodec = std::make_unique<ZarrV3CodecBytes>();
        else if (osName == ZarrV3CodecTranspose::NAME)
            poCodec = std::make_unique<ZarrV3CodecTranspose>();
        else if (osName == ZarrV3CodecCRC32C::NAME)
            poCodec = std::make_unique<ZarrV3CodecCRC32C>();
        else if (osName == ZarrV3CodecVLenUTF8::NAME)
            poCodec = std::make_unique<ZarrV3CodecVLenUTF8>();
        else if (osName == ZarrV3CodecShardingIndexed::NAME)
        {
            bShardingFound = true;
            poCodec = std::make_unique<ZarrV3CodecShardingIndexed>();
        }
        else
        {
            CPLError(CE_Failure, CPLE_NotSupported, "Unsupported codec: %s",
                     osName.c_str());
            return false;
        }

        if (poCodec->GetInputType() == ZarrV3Codec::IOType::ARRAY)
        {
            if (eLastType == ZarrV3Codec::IOType::BYTES)
            {
                CPLError(CE_Failure, CPLE_AppDefined,
                         "Cannot chain codec %s with %s",
                         poCodec->GetName().c_str(), osLastCodec.c_str());
                return false;
            }
        }
        else
        {
            InsertImplicitEndianCodecIfNeeded();
        }

        ZarrArrayMetadata oStepOutputArrayMetadata;
        if (osName == ZarrV3CodecShardingIndexed::NAME)
        {
            anBlockSizesBeforeSharding = oInputArrayMetadata.anBlockSizes;
        }
        if (!poCodec->InitFromConfiguration(oCodec["configuration"],
                                            oInputArrayMetadata,
                                            oStepOutputArrayMetadata,
                                            /* bEmitWarnings = */ true))
        {
            return false;
        }
        oInputArrayMetadata = std::move(oStepOutputArrayMetadata);
        eLastType = poCodec->GetOutputType();
        osLastCodec = poCodec->GetName();

        if (!poCodec->IsNoOp())
            m_apoCodecs.emplace_back(std::move(poCodec));
    }

    if (bShardingFound)
    {
        m_bPartialDecodingPossible =
            (m_apoCodecs.back()->GetName() == ZarrV3CodecShardingIndexed::NAME);
        if (!m_bPartialDecodingPossible)
        {
            m_bPartialDecodingPossible = false;
            // This is not an implementation limitation, but the result of a
            // badly thought dataset. Zarr-Python also emits a similar warning.
            CPLError(
                CE_Warning, CPLE_AppDefined,
                "Sharding codec found, but not in last position. Consequently "
                "partial shard decoding will not be possible");
            oInputArrayMetadata.anBlockSizes =
                std::move(anBlockSizesBeforeSharding);
        }
    }

    InsertImplicitEndianCodecIfNeeded();

    m_oCodecArray = oCodecs.Clone();
    oOutputArrayMetadata = std::move(oInputArrayMetadata);
    return true;
}

/************************************************************************/
/*                  ZarrV3CodecBytes::AllocateBuffer()                  */
/************************************************************************/

bool ZarrV3CodecSequence::AllocateBuffer(ZarrByteVectorQuickResize &abyBuffer,
                                         size_t nEltCount)
{
    if (!m_apoCodecs.empty())
    {
        const size_t nRawSize =
            nEltCount * m_oInputArrayMetadata.oElt.nativeSize;
        // Grow the temporary buffer a bit beyond the uncompressed size
        const size_t nMaxSize = nRawSize + nRawSize / 3 + 64;
        try
        {
            m_abyTmp.resize(nMaxSize);
        }
        catch (const std::exception &e)
        {
            CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
            return false;
        }
        m_abyTmp.resize(nRawSize);

        // Grow the input/output buffer too if we have several steps
        if (m_apoCodecs.size() >= 2 && abyBuffer.capacity() < nMaxSize)
        {
            const size_t nSize = abyBuffer.size();
            try
            {
                abyBuffer.resize(nMaxSize);
            }
            catch (const std::exception &e)
            {
                CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
                return false;
            }
            abyBuffer.resize(nSize);
        }
    }
    return true;
}

/************************************************************************/
/*                    ZarrV3CodecSequence::Encode()                     */
/************************************************************************/

bool ZarrV3CodecSequence::Encode(ZarrByteVectorQuickResize &abyBuffer)
{
    if (!AllocateBuffer(abyBuffer,
                        MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
        return false;
    for (const auto &poCodec : m_apoCodecs)
    {
        if (!poCodec->Encode(abyBuffer, m_abyTmp))
            return false;
        std::swap(abyBuffer, m_abyTmp);
    }
    return true;
}

/************************************************************************/
/*                    ZarrV3CodecSequence::Decode()                     */
/************************************************************************/

bool ZarrV3CodecSequence::Decode(ZarrByteVectorQuickResize &abyBuffer)
{
    if (!AllocateBuffer(abyBuffer,
                        MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
        return false;
    for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
    {
        const auto &poCodec = *iter;
        if (!poCodec->Decode(abyBuffer, m_abyTmp))
            return false;
        std::swap(abyBuffer, m_abyTmp);
    }
    return true;
}

/************************************************************************/
/*                 ZarrV3CodecSequence::DecodePartial()                 */
/************************************************************************/

bool ZarrV3CodecSequence::DecodePartial(VSIVirtualHandle *poFile,
                                        ZarrByteVectorQuickResize &abyBuffer,
                                        const std::vector<size_t> &anStartIdxIn,
                                        const std::vector<size_t> &anCountIn)
{
    CPLAssert(anStartIdxIn.size() == m_oInputArrayMetadata.anBlockSizes.size());
    CPLAssert(anStartIdxIn.size() == anCountIn.size());

    if (!AllocateBuffer(abyBuffer, MultiplyElements(anCountIn)))
        return false;

    // anStartIdxIn and anCountIn are expressed in the shape *before* encoding
    // We need to apply the potential transpositions before submitting them
    // to the decoder of the Array->Bytes decoder
    std::vector<size_t> anStartIdx(anStartIdxIn);
    std::vector<size_t> anCount(anCountIn);
    for (auto &poCodec : m_apoCodecs)
    {
        poCodec->ChangeArrayShapeForward(anStartIdx, anCount);
    }

    for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
    {
        const auto &poCodec = *iter;

        if (!poCodec->DecodePartial(poFile, abyBuffer, m_abyTmp, anStartIdx,
                                    anCount))
            return false;
        std::swap(abyBuffer, m_abyTmp);
    }
    return true;
}

/************************************************************************/
/*              ZarrV3CodecSequence::BatchDecodePartial()               */
/************************************************************************/

bool ZarrV3CodecSequence::BatchDecodePartial(
    VSIVirtualHandle *poFile, const char *pszFilename,
    const std::vector<std::pair<std::vector<size_t>, std::vector<size_t>>>
        &anRequests,
    std::vector<ZarrByteVectorQuickResize> &aResults)
{
    // Only batch-decode when sharding is the sole codec. If other codecs
    // (e.g. transpose) precede it, indices and output need codec-specific
    // transformations that BatchDecodePartial does not handle.
    if (m_apoCodecs.size() == 1)
    {
        auto *poSharding = dynamic_cast<ZarrV3CodecShardingIndexed *>(
            m_apoCodecs.back().get());
        if (poSharding)
        {
            return poSharding->BatchDecodePartial(poFile, pszFilename,
                                                  anRequests, aResults);
        }
    }

    // Fallback: sequential DecodePartial for non-sharding codec chains
    aResults.resize(anRequests.size());
    for (size_t i = 0; i < anRequests.size(); ++i)
    {
        if (!DecodePartial(poFile, aResults[i], anRequests[i].first,
                           anRequests[i].second))
            return false;
    }
    return true;
}

/************************************************************************/
/*             ZarrV3CodecSequence::GetInnerMostBlockSize()             */
/************************************************************************/

std::vector<size_t> ZarrV3CodecSequence::GetInnerMostBlockSize(
    const std::vector<size_t> &anOuterBlockSize) const
{
    auto chunkSize = anOuterBlockSize;
    for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
    {
        const auto &poCodec = *iter;
        if (m_bPartialDecodingPossible ||
            poCodec->GetName() != ZarrV3CodecShardingIndexed::NAME)
        {
            chunkSize = poCodec->GetInnerMostBlockSize(chunkSize);
        }
    }
    return chunkSize;
}
