/******************************************************************************
 *
 * Project:  OpenGIS Simple Features Reference Implementation
 * Purpose:  Helper to fill ArrowArray
 * Author:   Even Rouault <even dot rouault at spatialys.com>
 *
 ******************************************************************************
 * Copyright (c) 2022, Even Rouault <even dot rouault at spatialys.com>
 *
 * SPDX-License-Identifier: MIT
 ****************************************************************************/

#pragma once

//! @cond Doxygen_Suppress

#include <algorithm>
#include <limits>

#include "cpl_time.h"

#include "ogrsf_frmts.h"
#include "ogr_recordbatch.h"

class CPL_DLL OGRArrowArrayHelper
{
    OGRArrowArrayHelper(const OGRArrowArrayHelper &) = delete;
    OGRArrowArrayHelper &operator=(const OGRArrowArrayHelper &) = delete;

  public:
    bool m_bIncludeFID = false;
    int m_nMaxBatchSize = 0;
    int m_nChildren = 0;
    const int m_nFieldCount = 0;
    const int m_nGeomFieldCount = 0;
    std::vector<int> m_mapOGRFieldToArrowField{};
    std::vector<int> m_mapOGRGeomFieldToArrowField{};
    std::vector<bool> m_abNullableFields{};
    std::vector<uint32_t> m_anArrowFieldMaxAlloc{};
    std::vector<int> m_anTZFlags{};
    int64_t *m_panFIDValues = nullptr;
    struct ArrowArray *m_out_array = nullptr;

    static uint32_t GetMemLimit();

    static int
    GetMaxFeaturesInBatch(const CPLStringList &aosArrowArrayStreamOptions);

    OGRArrowArrayHelper(GDALDataset *poDS, OGRFeatureDefn *poFeatureDefn,
                        const CPLStringList &aosArrowArrayStreamOptions,
                        struct ArrowArray *out_array);

    bool SetNull(int iArrowField, int iFeat)
    {
        auto psArray = m_out_array->children[iArrowField];
        ++psArray->null_count;
        uint8_t *pabyNull =
            static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[0]));
        if (psArray->buffers[0] == nullptr)
        {
            pabyNull = static_cast<uint8_t *>(
                VSI_MALLOC_ALIGNED_AUTO_VERBOSE((m_nMaxBatchSize + 7) / 8));
            if (pabyNull == nullptr)
            {
                return false;
            }
            memset(pabyNull, 0xFF, (m_nMaxBatchSize + 7) / 8);
            psArray->buffers[0] = pabyNull;
        }
        pabyNull[iFeat / 8] &= static_cast<uint8_t>(~(1 << (iFeat % 8)));

        if (psArray->n_buffers == 3)
        {
            auto panOffsets =
                static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
            panOffsets[iFeat + 1] = panOffsets[iFeat];
        }
        return true;
    }

    inline static void SetBoolOn(struct ArrowArray *psArray, int iFeat)
    {
        static_cast<uint8_t *>(
            const_cast<void *>(psArray->buffers[1]))[iFeat / 8] |=
            static_cast<uint8_t>(1 << (iFeat % 8));
    }

    inline static void SetInt8(struct ArrowArray *psArray, int iFeat,
                               int8_t nVal)
    {
        static_cast<int8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    inline static void SetUInt8(struct ArrowArray *psArray, int iFeat,
                                uint8_t nVal)
    {
        static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    inline static void SetInt16(struct ArrowArray *psArray, int iFeat,
                                int16_t nVal)
    {
        static_cast<int16_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    inline static void SetUInt16(struct ArrowArray *psArray, int iFeat,
                                 uint16_t nVal)
    {
        static_cast<uint16_t *>(
            const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
    }

    inline static void SetInt32(struct ArrowArray *psArray, int iFeat,
                                int32_t nVal)
    {
        static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    inline static void SetUInt32(struct ArrowArray *psArray, int iFeat,
                                 uint32_t nVal)
    {
        static_cast<uint32_t *>(
            const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
    }

    inline static void SetInt64(struct ArrowArray *psArray, int iFeat,
                                int64_t nVal)
    {
        static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    inline static void SetUInt64(struct ArrowArray *psArray, int iFeat,
                                 uint64_t nVal)
    {
        static_cast<uint64_t *>(
            const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
    }

    inline static void SetFloat(struct ArrowArray *psArray, int iFeat,
                                float fVal)
    {
        static_cast<float *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            fVal;
    }

    inline static void SetDouble(struct ArrowArray *psArray, int iFeat,
                                 double dfVal)
    {
        static_cast<double *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            dfVal;
    }

    static void SetDate(struct ArrowArray *psArray, int iFeat,
                        struct tm &brokenDown, const OGRField &ogrField)
    {
        brokenDown.tm_year = ogrField.Date.Year - 1900;
        brokenDown.tm_mon = ogrField.Date.Month - 1;
        brokenDown.tm_mday = ogrField.Date.Day;
        brokenDown.tm_hour = 0;
        brokenDown.tm_min = 0;
        brokenDown.tm_sec = 0;
        static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            static_cast<int>(CPLYMDHMSToUnixTime(&brokenDown) / 86400);
    }

    static void SetDateTime(struct ArrowArray *psArray, int iFeat,
                            struct tm &brokenDown, int nFieldTZFlag,
                            const OGRField &ogrField)
    {
        brokenDown.tm_year = ogrField.Date.Year - 1900;
        brokenDown.tm_mon = ogrField.Date.Month - 1;
        brokenDown.tm_mday = ogrField.Date.Day;
        brokenDown.tm_hour = ogrField.Date.Hour;
        brokenDown.tm_min = ogrField.Date.Minute;
        brokenDown.tm_sec = static_cast<int>(ogrField.Date.Second);
        auto nVal =
            CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
            (static_cast<int>(ogrField.Date.Second * 1000 + 0.5f) % 1000);
        if (nFieldTZFlag >= OGR_TZFLAG_MIXED_TZ &&
            ogrField.Date.TZFlag > OGR_TZFLAG_MIXED_TZ)
        {
            // Convert for ogrField.Date.TZFlag to UTC
            const int TZOffset = (ogrField.Date.TZFlag - OGR_TZFLAG_UTC) * 15;
            const int TZOffsetMS = TZOffset * 60 * 1000;
            nVal -= TZOffsetMS;
        }
        static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
            nVal;
    }

    GByte *GetPtrForStringOrBinary(int iArrowField, int iFeat, size_t nLen)
    {
        auto psArray = m_out_array->children[iArrowField];
        auto panOffsets =
            static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
        const uint32_t nCurLength = static_cast<uint32_t>(panOffsets[iFeat]);
        if (nLen > m_anArrowFieldMaxAlloc[iArrowField] - nCurLength)
        {
            if (nLen >
                static_cast<uint32_t>(std::numeric_limits<int32_t>::max()) -
                    nCurLength)
            {
                CPLError(CE_Failure, CPLE_AppDefined,
                         "Too large string or binary content");
                return nullptr;
            }
            uint32_t nNewSize = nCurLength + static_cast<uint32_t>(nLen);
            if ((m_anArrowFieldMaxAlloc[iArrowField] >> 31) == 0)
            {
                const uint32_t nDoubleSize =
                    2U * m_anArrowFieldMaxAlloc[iArrowField];
                if (nNewSize < nDoubleSize)
                    nNewSize = nDoubleSize;
            }
            void *newBuffer = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(nNewSize);
            if (newBuffer == nullptr)
                return nullptr;
            m_anArrowFieldMaxAlloc[iArrowField] = nNewSize;
            memcpy(newBuffer, psArray->buffers[2], nCurLength);
            VSIFreeAligned(const_cast<void *>(psArray->buffers[2]));
            psArray->buffers[2] = newBuffer;
        }
        GByte *paby =
            static_cast<GByte *>(const_cast<void *>(psArray->buffers[2])) +
            nCurLength;
        panOffsets[iFeat + 1] = panOffsets[iFeat] + static_cast<int32_t>(nLen);
        return paby;
    }

    static void SetEmptyStringOrBinary(struct ArrowArray *psArray, int iFeat)
    {
        auto panOffsets =
            static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
        panOffsets[iFeat + 1] = panOffsets[iFeat];
    }

    void Shrink(int nFeatures)
    {
        if (nFeatures < m_nMaxBatchSize)
        {
            m_out_array->length = nFeatures;
            for (int i = 0; i < m_nChildren; i++)
            {
                m_out_array->children[i]->length = nFeatures;
            }
        }
    }

    void ClearArray()
    {
        if (m_out_array->release)
            m_out_array->release(m_out_array);
        memset(m_out_array, 0, sizeof(*m_out_array));
    }

    static bool FillDict(struct ArrowArray *psChild,
                         const OGRCodedFieldDomain *poCodedDomain);
};

//! @endcond