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

 *

 * Project:  GDAL

 * Purpose:  gdal "vector rasterize" subcommand

 * Author:   Alessandro Pasotti <elpaso at itopen dot it>

 *

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

 * Copyright (c) 2025, Alessandro Pasotti <elpaso at itopen dot it>

 *

 * SPDX-License-Identifier: MIT

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

#include <cmath>

#include "gdalalg_vector_rasterize.h"

#include "gdalalg_raster_write.h"

#include "cpl_conv.h"

#include "gdal_priv.h"

#include "gdal_utils.h"

//! @cond Doxygen_Suppress

#ifndef _

#define _(x) (x)

#endif

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

/*     GDALVectorRasterizeAlgorithm::GDALVectorRasterizeAlgorithm()     */

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

GDALVectorRasterizeAlgorithm::GDALVectorRasterizeAlgorithm(bool bStandaloneStep)

    : GDALPipelineStepAlgorithm(

          NAME, DESCRIPTION, HELP_URL,

          ConstructorOptions()

              .SetStandaloneStep(bStandaloneStep)

              .SetOutputFormatCreateCapability(GDAL_DCAP_CREATE))

{

    if (bStandaloneStep)

    {

        AddProgressArg();

        AddOutputFormatArg(&m_format)

            .AddMetadataItem(GAAMDI_REQUIRED_CAPABILITIES,

                             {GDAL_DCAP_RASTER, GDAL_DCAP_CREATE})

            .AddMetadataItem(GAAMDI_VRT_COMPATIBLE, {"false"});

        AddOpenOptionsArg(&m_openOptions);

        AddInputFormatsArg(&m_inputFormats)

            .AddMetadataItem(GAAMDI_REQUIRED_CAPABILITIES, {GDAL_DCAP_VECTOR});

        AddInputDatasetArg(&m_inputDataset, GDAL_OF_VECTOR)

            .SetMinCount(1)

            .SetMaxCount(1);

        AddOutputDatasetArg(&m_outputDataset, GDAL_OF_RASTER)

            .SetDatasetInputFlags(GADV_NAME | GADV_OBJECT);

        AddCreationOptionsArg(&m_creationOptions);

        AddOverwriteArg(&m_overwrite);

    }

    else

    {

        AddVectorHiddenInputDatasetArg();

    }

    AddBandArg(&m_bands, _("The band(s) to burn values into (1-based index)"));

    AddArg("invert", 0, _("Invert the rasterization"), &m_invert)

        .SetDefault(false);

    AddArg("all-touched", 0, _("Enables the ALL_TOUCHED rasterization option"),

           &m_allTouched);

    AddArg("burn", 0, _("Burn value"), &m_burnValues);

    AddArg("attribute-name", 'a', _("Attribute name"), &m_attributeName);

    AddArg("3d", 0,

           _("Indicates that a burn value should be extracted from the Z"

             " values of the feature"),

           &m_3d);

    AddLayerNameArg(&m_layerName).SetMutualExclusionGroup("layer-name-or-sql");

    AddArg("where", 0, _("SQL where clause"), &m_where);

    AddArg("sql", 0, _("SQL select statement"), &m_sql)

        .SetMutualExclusionGroup("layer-name-or-sql");

    AddArg("dialect", 0, _("SQL dialect"), &m_dialect);

    AddArg("nodata", 0, _("Assign a specified nodata value to output bands"),

           &m_nodata);

    AddArg("init", 0, _("Pre-initialize output bands with specified value"),

           &m_initValues);

    AddArg("like", 0,

           _("Dataset to use as a template for bounds, CRS and resolution"),

           &m_likeDataset, GDAL_OF_RASTER | GDAL_OF_VECTOR)

        .SetMetaVar("DATASET");

    AddArg("crs", 0, _("Override the projection for the output file"), &m_srs)

        .AddHiddenAlias("srs")

        .SetIsCRSArg(/*noneAllowed=*/false);

    AddArg("transformer-option", 0,

           _("Set a transformer option suitable to pass to "

             "GDALCreateGenImgProjTransformer2"),

           &m_transformerOption)

        .SetMetaVar("<NAME>=<VALUE>");

    AddArg("extent", 0, _("Set the target georeferenced extent"),

           &m_targetExtent)

        .SetMinCount(4)

        .SetMaxCount(4)

        .SetRepeatedArgAllowed(false)

        .SetMetaVar("<xmin>,<ymin>,<xmax>,<ymax>");

    auto &argResolution =

        AddArg("resolution", 0, _("Set the target resolution"),

               &m_targetResolution)

            .SetMinCount(2)

            .SetMaxCount(2)

            .SetRepeatedArgAllowed(false)

            .SetMetaVar("<xres>,<yres>")

            .SetMutualExclusionGroup("size-or-resolution");

    AddArg("target-aligned-pixels", 0,

           _("(target aligned pixels) Align the coordinates of the extent of "

             "the output file to the values of the resolution"),

           &m_tap)

        .AddAlias("tap")

        .AddDirectDependency(argResolution);

    AddArg("size", 0, _("Set the target size in pixels and lines"),

           &m_targetSize)

        .SetMinCount(2)

        .SetMaxCount(2)

        .SetRepeatedArgAllowed(false)

        .SetMetaVar("<xsize>,<ysize>")

        .SetMutualExclusionGroup("size-or-resolution");

    AddOutputDataTypeArg(&m_outputType);

    AddArg("optimization", 0,

           _("Force the algorithm used (results are identical)"),

           &m_optimization)

        .SetChoices("AUTO", "RASTER", "VECTOR")

        .SetDefault("AUTO");

    if (bStandaloneStep)

    {

        auto &addArg = AddArg("add", 0, _("Add to existing raster"), &m_add)

                           .SetDefault(false);

        auto &updateArg = AddUpdateArg(&m_update);

        addArg.AddValidationAction(

            [&updateArg]()

            {

                updateArg.Set(true);

                return true;

            });

    }

    AddValidationAction(

        [this]()

        {

            if (m_likeDataset.GetDatasetRef() &&

                (!m_targetExtent.empty() || !m_srs.empty()))

            {

                ReportError(

                    CE_Failure, CPLE_AppDefined,

                    "--like is mutually exclusive with --extent and --crs");

                return false;

            }

            return true;

        });

}

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

/*               GDALVectorRasterizeAlgorithm::RunStep()                */

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

bool GDALVectorRasterizeAlgorithm::RunStep(GDALPipelineStepRunContext &ctxt)

{

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

    CPLAssert(poSrcDS);

    CPLStringList aosOptions;

    if (m_bands.size())

    {

        for (int band : m_bands)

        {

            aosOptions.AddString("-b");

            aosOptions.AddString(CPLSPrintf("%d", band));

        }

    }

    if (m_invert)

    {

        aosOptions.AddString("-i");

    }

    if (m_allTouched)

    {

        aosOptions.AddString("-at");

    }

    if (m_burnValues.size())

    {

        for (double burnValue : m_burnValues)

        {

            aosOptions.AddString("-burn");

            aosOptions.AddString(CPLSPrintf("%.17g", burnValue));

        }

    }

    if (!m_attributeName.empty())

    {

        aosOptions.AddString("-a");

        aosOptions.AddString(m_attributeName.c_str());

    }

    if (m_3d)

    {

        aosOptions.AddString("-3d");

    }

    if (m_add)

    {

        aosOptions.AddString("-add");

        // Implies update

        m_update = true;

    }

    if (!m_layerName.empty())

    {

        aosOptions.AddString("-l");

        aosOptions.AddString(m_layerName.c_str());

    }

    if (!m_where.empty())

    {

        aosOptions.AddString("-where");

        aosOptions.AddString(m_where.c_str());

    }

    if (!m_sql.empty())

    {

        aosOptions.AddString("-sql");

        aosOptions.AddString(m_sql.c_str());

    }

    if (!m_dialect.empty())

    {

        aosOptions.AddString("-dialect");

        aosOptions.AddString(m_dialect.c_str());

    }

    std::string outputFilename;

    if (m_standaloneStep)

    {

        outputFilename = m_outputDataset.GetName();

        if (!m_format.empty())

        {

            aosOptions.AddString("-of");

            aosOptions.AddString(m_format.c_str());

        }

        for (const std::string &co : m_creationOptions)

        {

            aosOptions.AddString("-co");

            aosOptions.AddString(co.c_str());

        }

    }

    else

    {

        outputFilename = CPLGenerateTempFilenameSafe("_rasterize.tif");

        aosOptions.AddString("-of");

        aosOptions.AddString("GTiff");

        aosOptions.AddString("-co");

        aosOptions.AddString("TILED=YES");

    }

    if (!std::isnan(m_nodata))

    {

        if (m_update)

        {

            ReportError(

                CE_Failure, CPLE_AppDefined,

                "Cannot specify --nodata when updating an existing raster.");

            return false;

        }

        aosOptions.AddString("-a_nodata");

        aosOptions.AddString(CPLSPrintf("%.17g", m_nodata));

    }

    if (m_initValues.size())

    {

        for (double initValue : m_initValues)

        {

            aosOptions.AddString("-init");

            aosOptions.AddString(CPLSPrintf("%.17g", initValue));

        }

    }

    if (auto *poLikeDS = m_likeDataset.GetDatasetRef())

    {

        OGREnvelope sExtent;

        if (poLikeDS->GetExtent(&sExtent) != CE_None)

        {

            ReportError(CE_Failure, CPLE_AppDefined,

                        "Cannot get extent of '%s'",

                        poLikeDS->GetDescription());

            return false;

        }

        aosOptions.AddString("-te");

        aosOptions.AddString(CPLSPrintf("%.17g", sExtent.MinX));

        aosOptions.AddString(CPLSPrintf("%.17g", sExtent.MinY));

        aosOptions.AddString(CPLSPrintf("%.17g", sExtent.MaxX));

        aosOptions.AddString(CPLSPrintf("%.17g", sExtent.MaxY));

        const auto poSRS = poLikeDS->GetSpatialRef();

        if (poSRS)

        {

            const char *const options[] = {"FORMAT=WKT2_2019", nullptr};

            const std::string osCRS = poSRS->exportToWkt(options);

            aosOptions.AddString("-a_srs");

            aosOptions.AddString(osCRS.c_str());

        }

        if (m_targetResolution.empty() && m_targetSize.empty())

        {

            GDALGeoTransform gt;

            if (poLikeDS->GetGeoTransform(gt) == CE_None)

            {

                if (gt.xrot != 0 || gt.yrot != 0)

                {

                    ReportError(CE_Failure, CPLE_NotSupported,

                                "Geotransform matrix of '%s' has non-zero "

                                "rotation terms",

                                poLikeDS->GetDescription());

                    return false;

                }

                aosOptions.AddString("-ts");

                aosOptions.AddString(

                    CPLSPrintf("%d", poLikeDS->GetRasterXSize()));

                aosOptions.AddString(

                    CPLSPrintf("%d", poLikeDS->GetRasterYSize()));

            }

        }

    }

    else

    {

        if (m_targetExtent.size())

        {

            aosOptions.AddString("-te");

            for (double targetExtent : m_targetExtent)

            {

                aosOptions.AddString(CPLSPrintf("%.17g", targetExtent));

            }

        }

        if (!m_srs.empty())

        {

            if (m_update)

            {

                ReportError(

                    CE_Failure, CPLE_AppDefined,

                    "Cannot specify --crs when updating an existing raster.");

                return false;

            }

            aosOptions.AddString("-a_srs");

            aosOptions.AddString(m_srs.c_str());

        }

    }

    if (m_transformerOption.size())

    {

        for (const auto &to : m_transformerOption)

        {

            aosOptions.AddString("-to");

            aosOptions.AddString(to.c_str());

        }

    }

    if (m_tap)

    {

        aosOptions.AddString("-tap");

    }

    if (m_targetResolution.size())

    {

        if (m_update)

        {

            ReportError(CE_Failure, CPLE_AppDefined,

                        "Cannot specify --resolution when updating an existing "

                        "raster.");

            return false;

        }

        aosOptions.AddString("-tr");

        for (double targetResolution : m_targetResolution)

        {

            aosOptions.AddString(CPLSPrintf("%.17g", targetResolution));

        }

    }

    else if (m_targetSize.size())

    {

        if (m_update)

        {

            ReportError(

                CE_Failure, CPLE_AppDefined,

                "Cannot specify --size when updating an existing raster.");

            return false;

        }

        aosOptions.AddString("-ts");

        for (int targetSize : m_targetSize)

        {

            aosOptions.AddString(CPLSPrintf("%d", targetSize));

        }

    }

    else if (!m_likeDataset.GetDatasetRef() && !m_outputDataset.GetDatasetRef())

    {

        ReportError(

            CE_Failure, CPLE_AppDefined,

            "Must specify output resolution (--resolution) or size (--size) "

            "when writing rasterized features to a new dataset.");

        return false;

    }

    if (!m_outputType.empty())

    {

        if (m_update)

        {

            ReportError(CE_Failure, CPLE_AppDefined,

                        "Cannot specify --output-data-type when updating an "

                        "existing raster.");

            return false;

        }

        aosOptions.AddString("-ot");

        aosOptions.AddString(m_outputType.c_str());

    }

    if (!m_optimization.empty())

    {

        aosOptions.AddString("-optim");

        aosOptions.AddString(m_optimization.c_str());

    }

    bool bOK = false;

    std::unique_ptr<GDALRasterizeOptions, decltype(&GDALRasterizeOptionsFree)>

        psOptions{GDALRasterizeOptionsNew(aosOptions.List(), nullptr),

                  GDALRasterizeOptionsFree};

    if (psOptions)

    {

        GDALRasterizeOptionsSetProgress(psOptions.get(), ctxt.m_pfnProgress,

                                        ctxt.m_pProgressData);

        GDALDatasetH hDstDS =

            GDALDataset::ToHandle(m_outputDataset.GetDatasetRef());

        GDALDatasetH hSrcDS = GDALDataset::ToHandle(poSrcDS);

        auto poRetDS = GDALDataset::FromHandle(GDALRasterize(

            outputFilename.c_str(), hDstDS, hSrcDS, psOptions.get(), nullptr));

        bOK = poRetDS != nullptr;

        if (!hDstDS)

        {

            if (!m_standaloneStep && poRetDS)

            {

                VSIUnlink(outputFilename.c_str());

                poRetDS->MarkSuppressOnClose();

            }

            m_outputDataset.Set(std::unique_ptr<GDALDataset>(poRetDS));

        }

    }

    return bOK;

}

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

/*               GDALVectorRasterizeAlgorithm::RunImpl()                */

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

bool GDALVectorRasterizeAlgorithm::RunImpl(GDALProgressFunc pfnProgress,

                                           void *pProgressData)

{

    GDALPipelineStepRunContext stepCtxt;

    stepCtxt.m_pfnProgress = pfnProgress;

    stepCtxt.m_pProgressData = pProgressData;

    return RunPreStepPipelineValidations() && RunStep(stepCtxt);

}

GDALVectorRasterizeAlgorithmStandalone::

    ~GDALVectorRasterizeAlgorithmStandalone() = default;

//! @endcond