/******************************************************************************
 *
 * Project:  GDAL
 * Purpose:  gdal "raster rgb-to-palette" subcommand
 * Author:   Even Rouault <even dot rouault at spatialys.com>
 *
 ******************************************************************************
 * Copyright (c) 2025, Even Rouault <even dot rouault at spatialys.com>
 *
 * SPDX-License-Identifier: MIT
 ****************************************************************************/

#include "gdalalg_raster_rgb_to_palette.h"

#include "cpl_string.h"
#include "gdal_alg.h"
#include "gdal_alg_priv.h"
#include "gdal_priv.h"

#include <algorithm>
#include <limits>

//! @cond Doxygen_Suppress

#ifndef _
#define _(x) (x)
#endif

/************************************************************************/
/*                  GDALRasterRGBToPaletteAlgorithm()                   */
/************************************************************************/

GDALRasterRGBToPaletteAlgorithm::GDALRasterRGBToPaletteAlgorithm(
    bool standaloneStep)
    : GDALRasterPipelineNonNativelyStreamingAlgorithm(NAME, DESCRIPTION,
                                                      HELP_URL, standaloneStep)
{
    AddArg("color-count", 0,
           _("Select the number of colors in the generated color table"),
           &m_colorCount)
        .SetDefault(m_colorCount)
        .SetMinValueIncluded(2)
        .SetMaxValueIncluded(256);
    AddArg("color-map", 0, _("Color map filename"), &m_colorMap);
    AddArg("output-nodata", 0, _("Output nodata value"), &m_dstNoData)
        .AddHiddenAlias("dst-nodata")
        .SetMinValueIncluded(0)
        .SetMaxValueIncluded(255);
    AddArg("no-dither", 0, _("Disable Floyd-Steinberg dithering"), &m_noDither);
    AddArg("bit-depth", 0,
           _("Bit depth of color palette component (8 bit causes longer "
             "computation time)"),
           &m_bitDepth)
        .SetDefault(m_bitDepth)
        .SetChoices("5", "8");
}

/************************************************************************/
/*              GDALRasterRGBToPaletteAlgorithm::RunStep()              */
/************************************************************************/

bool GDALRasterRGBToPaletteAlgorithm::RunStep(GDALPipelineStepRunContext &ctxt)
{
    auto pfnProgress = ctxt.m_pfnProgress;
    auto pProgressData = ctxt.m_pProgressData;
    auto poSrcDS = m_inputDataset[0].GetDatasetRef();
    CPLAssert(poSrcDS);

    const int nSrcBandCount = poSrcDS->GetRasterCount();
    if (nSrcBandCount < 3)
    {
        ReportError(CE_Failure, CPLE_NotSupported,
                    "Input dataset must have at least 3 bands");
        return false;
    }
    else if (nSrcBandCount == 4 &&
             poSrcDS->GetRasterBand(4)->GetColorInterpretation() ==
                 GCI_AlphaBand)
    {
        // nothing to do
    }
    else if (nSrcBandCount >= 4)
    {
        ReportError(
            CE_Warning, CPLE_AppDefined,
            "Only R,G,B bands of input dataset will be taken into account");
    }

    std::map<GDALColorInterp, GDALRasterBandH> mapBands;
    int nFound = 0;
    for (int i = 1; i <= nSrcBandCount; ++i)
    {
        auto poSrcBand = poSrcDS->GetRasterBand(i);
        if (poSrcBand->GetRasterDataType() != GDT_UInt8)
        {
            ReportError(CE_Failure, CPLE_NotSupported,
                        "Non-byte band found and not supported");
            return false;
        }
        const auto eColorInterp = poSrcBand->GetColorInterpretation();
        for (const auto eInterestColorInterp :
             {GCI_RedBand, GCI_GreenBand, GCI_BlueBand})
        {
            if (eColorInterp == eInterestColorInterp)
            {
                if (mapBands.find(eColorInterp) != mapBands.end())
                {
                    ReportError(CE_Failure, CPLE_NotSupported,
                                "Several %s bands found",
                                GDALGetColorInterpretationName(eColorInterp));
                    return false;
                }
                ++nFound;
                mapBands[eColorInterp] = GDALRasterBand::ToHandle(poSrcBand);
            }
        }
    }

    if (nFound < 3)
    {
        if (nFound > 0)
        {
            ReportError(
                CE_Warning, CPLE_AppDefined,
                "Assuming first band is red, second green and third blue, "
                "despite at least one band with one of those color "
                "interpretation "
                "found");
        }
        mapBands[GCI_RedBand] =
            GDALRasterBand::ToHandle(poSrcDS->GetRasterBand(1));
        mapBands[GCI_GreenBand] =
            GDALRasterBand::ToHandle(poSrcDS->GetRasterBand(2));
        mapBands[GCI_BlueBand] =
            GDALRasterBand::ToHandle(poSrcDS->GetRasterBand(3));
    }

    auto poTmpDS = CreateTemporaryDataset(
        poSrcDS->GetRasterXSize(), poSrcDS->GetRasterYSize(), 1, GDT_UInt8,
        /* bTiledIfPossible = */ true, poSrcDS, /* bCopyMetadata = */ true);
    if (!poTmpDS)
        return false;

    const double oneOverStep = 1.0 / ((m_colorMap.empty() ? 1 : 0) + 1);

    if (m_colorMap.empty() && m_dstNoData < 0)
    {
        int bSrcHasNoDataR = FALSE;
        const double dfSrcNoDataR =
            GDALGetRasterNoDataValue(mapBands[GCI_RedBand], &bSrcHasNoDataR);
        int bSrcHasNoDataG = FALSE;
        const double dfSrcNoDataG =
            GDALGetRasterNoDataValue(mapBands[GCI_GreenBand], &bSrcHasNoDataG);
        int bSrcHasNoDataB = FALSE;
        const double dfSrcNoDataB =
            GDALGetRasterNoDataValue(mapBands[GCI_BlueBand], &bSrcHasNoDataB);
        if (bSrcHasNoDataR && bSrcHasNoDataG && bSrcHasNoDataB &&
            dfSrcNoDataR == dfSrcNoDataG && dfSrcNoDataR == dfSrcNoDataB &&
            dfSrcNoDataR >= 0 && dfSrcNoDataR <= 255 &&
            std::round(dfSrcNoDataR) == dfSrcNoDataR)
        {
            m_dstNoData = 0;
        }
        else
        {
            const int nMaskFlags = GDALGetMaskFlags(mapBands[GCI_RedBand]);
            if ((nMaskFlags & GMF_PER_DATASET))
                m_dstNoData = 0;
        }
    }

    GDALColorTable oCT;

    bool bOK = true;
    double dfLastProgress = 0;
    std::unique_ptr<void, decltype(&GDALDestroyScaledProgress)> pScaledData(
        nullptr, GDALDestroyScaledProgress);
    if (m_colorMap.empty())
    {
        pScaledData.reset(GDALCreateScaledProgress(0, oneOverStep, pfnProgress,
                                                   pProgressData));
        dfLastProgress = oneOverStep;

        const int nXSize = poSrcDS->GetRasterXSize();
        const int nYSize = poSrcDS->GetRasterYSize();

        if (m_dstNoData >= 0 && m_colorCount == 256)
            --m_colorCount;
        if (nYSize == 0)
        {
            bOK = false;
        }
        else if (static_cast<GUInt32>(nXSize) <
                 std::numeric_limits<GUInt32>::max() /
                     static_cast<GUInt32>(nYSize))
        {
            bOK =
                GDALComputeMedianCutPCTInternal(
                    mapBands[GCI_RedBand], mapBands[GCI_GreenBand],
                    mapBands[GCI_BlueBand], nullptr, nullptr, nullptr, nullptr,
                    m_colorCount, m_bitDepth, static_cast<GUInt32 *>(nullptr),
                    GDALColorTable::ToHandle(&oCT),
                    pScaledData ? GDALScaledProgress : nullptr,
                    pScaledData.get()) == CE_None;
        }
        else
        {
            bOK =
                GDALComputeMedianCutPCTInternal(
                    mapBands[GCI_RedBand], mapBands[GCI_GreenBand],
                    mapBands[GCI_BlueBand], nullptr, nullptr, nullptr, nullptr,
                    m_colorCount, m_bitDepth, static_cast<GUIntBig *>(nullptr),
                    GDALColorTable::ToHandle(&oCT),
                    pScaledData ? GDALScaledProgress : nullptr,
                    pScaledData.get()) == CE_None;
        }
    }
    else
    {
        GDALDriverH hDriver;
        if ((hDriver = GDALIdentifyDriver(m_colorMap.c_str(), nullptr)) !=
                nullptr &&
            // Palette .txt files may be misidentified by the XYZ driver
            !EQUAL(GDALGetDescription(hDriver), "XYZ"))
        {
            auto poPaletteDS = std::unique_ptr<GDALDataset>(GDALDataset::Open(
                m_colorMap.c_str(), GDAL_OF_RASTER | GDAL_OF_VERBOSE_ERROR,
                nullptr, nullptr, nullptr));
            bOK = poPaletteDS != nullptr && poPaletteDS->GetRasterCount() > 0;
            if (bOK)
            {
                const auto poCT =
                    poPaletteDS->GetRasterBand(1)->GetColorTable();
                if (poCT)
                {
                    oCT = *poCT;
                }
                else
                {
                    bOK = false;
                    ReportError(CE_Failure, CPLE_AppDefined,
                                "Dataset '%s' does not contain a color table",
                                m_colorMap.c_str());
                }
            }
        }
        else
        {
            auto poCT = GDALColorTable::LoadFromFile(m_colorMap.c_str());
            bOK = poCT != nullptr;
            if (bOK)
            {
                oCT = std::move(*(poCT.get()));
            }
        }
    }

    m_colorCount = oCT.GetColorEntryCount();

    if (m_dstNoData >= 0)
    {
        for (int i = std::min(255, m_colorCount); i > m_dstNoData; --i)
        {
            // Create a temporary copy to avoid use after free when SetColorEntry()
            // resizes the underlying vector.
            const GDALColorEntry sEntry = *(oCT.GetColorEntry(i - 1));
            oCT.SetColorEntry(i, &sEntry);
        }

        poTmpDS->GetRasterBand(1)->SetNoDataValue(m_dstNoData);
        GDALColorEntry sEntry = {0, 0, 0, 0};
        oCT.SetColorEntry(m_dstNoData, &sEntry);
    }

    if (bOK)
    {
        poTmpDS->GetRasterBand(1)->SetColorTable(&oCT);

        pScaledData.reset(GDALCreateScaledProgress(dfLastProgress, 1.0,
                                                   pfnProgress, pProgressData));

        bOK = GDALDitherRGB2PCTInternal(
                  mapBands[GCI_RedBand], mapBands[GCI_GreenBand],
                  mapBands[GCI_BlueBand],
                  GDALRasterBand::ToHandle(poTmpDS->GetRasterBand(1)),
                  GDALColorTable::ToHandle(&oCT), m_bitDepth,
                  /* pasDynamicColorMap = */ nullptr, !m_noDither,
                  pScaledData ? GDALScaledProgress : nullptr,
                  pScaledData.get()) == CE_None;
    }

    if (bOK)
    {
        m_outputDataset.Set(std::move(poTmpDS));
        if (pfnProgress)
            pfnProgress(1.0, "", pProgressData);
    }

    return bOK;
}

GDALRasterRGBToPaletteAlgorithmStandalone::
    ~GDALRasterRGBToPaletteAlgorithmStandalone() = default;

//! @endcond
