/******************************************************************************

 *

 * Project:  GDAL

 * Purpose:  "reclassify" step of "raster pipeline"

 * Author:   Daniel Baston

 *

 ******************************************************************************

 * Copyright (c) 2025, ISciences LLC

 *

 * SPDX-License-Identifier: MIT

 ****************************************************************************/

#include "gdalalg_raster_reclassify.h"

#include "cpl_vsi_virtual.h"

#include "gdal_priv.h"

#include "gdal_utils.h"

#include "../frmts/vrt/vrtdataset.h"

#include "../frmts/vrt/vrtreclassifier.h"

#include <array>

//! @cond Doxygen_Suppress

#ifndef _

#define _(x) (x)

#endif

/************************************************************************/

/*    GDALRasterReclassifyAlgorithm::GDALRasterReclassifyAlgorithm()    */

/************************************************************************/

GDALRasterReclassifyAlgorithm::GDALRasterReclassifyAlgorithm(

    bool standaloneStep)

    : GDALRasterPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,

                                      standaloneStep)

{

    AddArg("mapping", 'm',

           _("Reclassification mappings (or specify a @<filename> to point to "

             "a file containing mappings"),

           &m_mapping)

        .SetRequired();

    AddArg("keep-color-table", 0, _("Preserve the input color table"),

           &m_keepColorTable);

    AddOutputDataTypeArg(&m_type);

}

/************************************************************************/

/*                 GDALRasterReclassifyValidateMappings                 */

/************************************************************************/

static bool GDALReclassifyValidateMappings(GDALDataset &input,

                                           const std::string &mappings,

                                           GDALDataType eDstType)

{

    int hasNoData;

    std::optional<double> noData =

        input.GetRasterBand(1)->GetNoDataValue(&hasNoData);

    if (!hasNoData)

    {

        noData.reset();

    }

    gdal::Reclassifier reclassifier;

    return reclassifier.Init(mappings.c_str(), noData, eDstType) == CE_None;

}

/************************************************************************/

/*                 GDALRasterReclassifyCreateVRTDerived                 */

/************************************************************************/

static std::unique_ptr<GDALDataset>

GDALReclassifyCreateVRTDerived(GDALDataset &input, const std::string &mappings,

                               GDALDataType eDstType, bool keepColorTable)

{

    const auto nX = input.GetRasterXSize();

    const auto nY = input.GetRasterYSize();

    auto poDS = VRTDataset::CreateVRTDataset("", nX, nY, 0, eDstType, nullptr);

    GDALGeoTransform gt;

    if (input.GetGeoTransform(gt) == CE_None)

        poDS->SetGeoTransform(gt);

    poDS->SetSpatialRef(input.GetSpatialRef());

    CPLStringList papszBandArgs;

    papszBandArgs.SetNameValue("subclass", "VRTDerivedRasterBand");

    for (int iBand = 1; iBand <= input.GetRasterCount(); ++iBand)

    {

        GDALRasterBand *poSrcBand = input.GetRasterBand(iBand);

        const GDALDataType srcType = poSrcBand->GetRasterDataType();

        const GDALDataType bandType =

            eDstType == GDT_Unknown ? srcType : eDstType;

        const GDALDataType xferType = GDALDataTypeUnion(srcType, bandType);

        if (poDS->AddBand(bandType, papszBandArgs) != CE_None)

        {

            return nullptr;

        }

        VRTDerivedRasterBand *poDstBand =

            cpl::down_cast<VRTDerivedRasterBand *>(poDS->GetRasterBand(iBand));

        poDstBand->SetSourceTransferType(xferType);

        poDstBand->SetPixelFunctionName("reclassify");

        poDstBand->AddPixelFunctionArgument("mapping", mappings.c_str());

        if (keepColorTable && poSrcBand->GetColorTable() != nullptr)

        {

            poDstBand->SetColorTable(poSrcBand->GetColorTable());

        }

        GDALCopyNoDataValue(poDstBand, poSrcBand);

        poDstBand->AddSimpleSource(poSrcBand);

    }

    return poDS;

}

/************************************************************************/

/*               GDALRasterReclassifyAlgorithm::RunStep()               */

/************************************************************************/

bool GDALRasterReclassifyAlgorithm::RunStep(GDALPipelineStepRunContext &)

{

    const auto poSrcDS = m_inputDataset[0].GetDatasetRef();

    CPLAssert(poSrcDS);

    CPLAssert(m_outputDataset.GetName().empty());

    CPLAssert(!m_outputDataset.GetDatasetRef());

    // Already validated by argument parser

    const GDALDataType eDstType =

        m_type.empty() ? GDT_Unknown : GDALGetDataTypeByName(m_type.c_str());

    const auto nErrorCount = CPLGetErrorCounter();

    if (!m_mapping.empty() && m_mapping[0] == '@')

    {

        auto f =

            VSIVirtualHandleUniquePtr(VSIFOpenL(m_mapping.c_str() + 1, "r"));

        if (!f)

        {

            ReportError(CE_Failure, CPLE_FileIO, "Cannot open %s",

                        m_mapping.c_str() + 1);

            return false;

        }

        m_mapping.clear();

        try

        {

            constexpr int MAX_CHARS_PER_LINE = 1000 * 1000;

            constexpr size_t MAX_MAPPING_SIZE = 10 * 1000 * 1000;

            while (const char *line =

                       CPLReadLine2L(f.get(), MAX_CHARS_PER_LINE, nullptr))

            {

                while (isspace(*line))

                {

                    line++;

                }

                if (line[0])

                {

                    if (!m_mapping.empty())

                    {

                        m_mapping.append(";");

                    }

                    const char *comment = strchr(line, '#');

                    if (!comment)

                    {

                        m_mapping.append(line);

                    }

                    else

                    {

                        m_mapping.append(line,

                                         static_cast<size_t>(comment - line));

                    }

                    if (m_mapping.size() > MAX_MAPPING_SIZE)

                    {

                        ReportError(CE_Failure, CPLE_AppDefined,

                                    "Too large mapping size");

                        return false;

                    }

                }

            }

        }

        catch (const std::exception &)

        {

            ReportError(CE_Failure, CPLE_OutOfMemory,

                        "Out of memory while ingesting mapping file");

        }

    }

    if (nErrorCount == CPLGetErrorCounter())

    {

        if (!GDALReclassifyValidateMappings(*poSrcDS, m_mapping, eDstType))

        {

            return false;

        }

        m_outputDataset.Set(GDALReclassifyCreateVRTDerived(

            *poSrcDS, m_mapping, eDstType, m_keepColorTable));

    }

    return m_outputDataset.GetDatasetRef() != nullptr;

}

GDALRasterReclassifyAlgorithmStandalone::

    ~GDALRasterReclassifyAlgorithmStandalone() = default;

//! @endcond