/******************************************************************************
 *
 * Project:  GDAL
 * Purpose:  "sql" step of "vector pipeline"
 * 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_vector_sql.h"

#include "gdal_priv.h"
#include "ogrsf_frmts.h"
#include "ogrlayerpool.h"

#include <mutex>
#include <set>

//! @cond Doxygen_Suppress

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

/************************************************************************/
/*           GDALVectorSQLAlgorithm::GetConstructorOptions()            */
/************************************************************************/

/* static */ GDALVectorSQLAlgorithm::ConstructorOptions
GDALVectorSQLAlgorithm::GetConstructorOptions(bool standaloneStep)
{
    ConstructorOptions opts;
    opts.SetStandaloneStep(standaloneStep);
    opts.SetOutputDatasetRequired(false);
    opts.SetAddInputLayerNameArgument(false);
    opts.SetAddOutputLayerNameArgument(false);
    return opts;
}

/************************************************************************/
/*           GDALVectorSQLAlgorithm::GDALVectorSQLAlgorithm()           */
/************************************************************************/

GDALVectorSQLAlgorithm::GDALVectorSQLAlgorithm(bool standaloneStep)
    : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
                                      GetConstructorOptions(standaloneStep))
{
    auto &sqlArg = AddArg("sql", 0, _("SQL statement(s)"), &m_sql)
                       .SetRequired()
                       .SetPackedValuesAllowed(false)
                       .SetReadFromFileAtSyntaxAllowed()
                       .SetMetaVar("<statement>|@<filename>")
                       .SetRemoveSQLCommentsEnabled();
    if (!standaloneStep)
        sqlArg.SetPositional();
    AddArg(GDAL_ARG_NAME_OUTPUT_LAYER, standaloneStep ? 0 : 'l',
           _("Output layer name(s)"), &m_outputLayer);
    AddArg("dialect", 0, _("SQL dialect (e.g. OGRSQL, SQLITE)"), &m_dialect);
}

/************************************************************************/
/*                   GDALVectorSQLAlgorithmDataset                      */
/************************************************************************/

namespace
{
class GDALVectorSQLAlgorithmDataset final : public GDALDataset
{
    GDALDataset &m_oSrcDS;
    std::vector<OGRLayer *> m_layers{};

    CPL_DISALLOW_COPY_ASSIGN(GDALVectorSQLAlgorithmDataset)

  public:
    explicit GDALVectorSQLAlgorithmDataset(GDALDataset &oSrcDS)
        : m_oSrcDS(oSrcDS)
    {
        m_oSrcDS.Reference();
    }

    ~GDALVectorSQLAlgorithmDataset() override
    {
        for (OGRLayer *poLayer : m_layers)
            m_oSrcDS.ReleaseResultSet(poLayer);
        m_oSrcDS.ReleaseRef();
    }

    void AddLayer(OGRLayer *poLayer)
    {
        m_layers.push_back(poLayer);
    }

    int GetLayerCount() const override
    {
        return static_cast<int>(m_layers.size());
    }

    OGRLayer *GetLayer(int idx) const override
    {
        return idx >= 0 && idx < GetLayerCount() ? m_layers[idx] : nullptr;
    }
};
}  // namespace

/************************************************************************/
/*               GDALVectorSQLAlgorithmDatasetMultiLayer                */
/************************************************************************/

namespace
{

class ProxiedSQLLayer final : public OGRProxiedLayer
{
    mutable OGRFeatureDefn *m_poLayerDefn = nullptr;
    mutable std::mutex m_oMutex{};

    CPL_DISALLOW_COPY_ASSIGN(ProxiedSQLLayer)

  public:
    ProxiedSQLLayer(const std::string &osName, OGRLayerPool *poPoolIn,
                    OpenLayerFunc pfnOpenLayerIn,
                    ReleaseLayerFunc pfnReleaseLayerIn,
                    FreeUserDataFunc pfnFreeUserDataIn, void *pUserDataIn)
        : OGRProxiedLayer(poPoolIn, pfnOpenLayerIn, pfnReleaseLayerIn,
                          pfnFreeUserDataIn, pUserDataIn)
    {
        SetDescription(osName.c_str());
    }

    ~ProxiedSQLLayer() override
    {
        if (m_poLayerDefn)
            m_poLayerDefn->Release();
    }

    const char *GetName() const override
    {
        return GetDescription();
    }

    const OGRFeatureDefn *GetLayerDefn() const override
    {
        std::lock_guard oLock(m_oMutex);

        if (!m_poLayerDefn)
        {
            m_poLayerDefn = OGRProxiedLayer::GetLayerDefn()->Clone();
            m_poLayerDefn->SetName(GetDescription());
        }
        return m_poLayerDefn;
    }
};

class GDALVectorSQLAlgorithmDatasetMultiLayer final : public GDALDataset
{
    // We can't safely have 2 SQL layers active simultaneously on the same
    // source dataset. So each time we access one, we must close the last
    // active one.
    OGRLayerPool m_oPool{1};
    GDALDataset &m_oSrcDS;
    std::vector<std::unique_ptr<ProxiedSQLLayer>> m_layers{};

    struct UserData
    {
        GDALDataset &oSrcDS;
        std::string osSQL{};
        std::string osDialect{};
        std::string osLayerName{};

        UserData(GDALDataset &oSrcDSIn, const std::string &osSQLIn,
                 const std::string &osDialectIn,
                 const std::string &osLayerNameIn)
            : oSrcDS(oSrcDSIn), osSQL(osSQLIn), osDialect(osDialectIn),
              osLayerName(osLayerNameIn)
        {
        }
        CPL_DISALLOW_COPY_ASSIGN(UserData)
    };

    CPL_DISALLOW_COPY_ASSIGN(GDALVectorSQLAlgorithmDatasetMultiLayer)

  public:
    explicit GDALVectorSQLAlgorithmDatasetMultiLayer(GDALDataset &oSrcDS)
        : m_oSrcDS(oSrcDS)
    {
        m_oSrcDS.Reference();
    }

    ~GDALVectorSQLAlgorithmDatasetMultiLayer() override
    {
        m_layers.clear();
        m_oSrcDS.ReleaseRef();
    }

    void AddLayer(const std::string &osSQL, const std::string &osDialect,
                  const std::string &osLayerName)
    {
        const auto OpenLayer = [](void *pUserDataIn)
        {
            UserData *pUserData = static_cast<UserData *>(pUserDataIn);
            return pUserData->oSrcDS.ExecuteSQL(
                pUserData->osSQL.c_str(), nullptr,
                pUserData->osDialect.empty() ? nullptr
                                             : pUserData->osDialect.c_str());
        };

        const auto CloseLayer = [](OGRLayer *poLayer, void *pUserDataIn)
        {
            UserData *pUserData = static_cast<UserData *>(pUserDataIn);
            pUserData->oSrcDS.ReleaseResultSet(poLayer);
        };

        const auto DeleteUserData = [](void *pUserDataIn)
        { delete static_cast<UserData *>(pUserDataIn); };

        auto pUserData = new UserData(m_oSrcDS, osSQL, osDialect, osLayerName);
        m_layers.emplace_back(std::make_unique<ProxiedSQLLayer>(
            osLayerName, &m_oPool, OpenLayer, CloseLayer, DeleteUserData,
            pUserData));
    }

    int GetLayerCount() const override
    {
        return static_cast<int>(m_layers.size());
    }

    OGRLayer *GetLayer(int idx) const override
    {
        return idx >= 0 && idx < GetLayerCount() ? m_layers[idx].get()
                                                 : nullptr;
    }
};
}  // namespace

/************************************************************************/
/*                 GDALVectorSQLAlgorithm::RunStep()                    */
/************************************************************************/

bool GDALVectorSQLAlgorithm::RunStep(GDALPipelineStepRunContext &)
{
    auto poSrcDS = m_inputDataset[0].GetDatasetRef();
    CPLAssert(poSrcDS);

    auto outputArg = GetArg(GDAL_ARG_NAME_OUTPUT);
    if (outputArg && !outputArg->IsExplicitlySet())
    {
        // Mode where we update a dataset.
        for (const auto &sql : m_sql)
        {
            const auto nErrorCounter = CPLGetErrorCounter();
            OGRLayer *poLayer = poSrcDS->ExecuteSQL(
                sql.c_str(), nullptr,
                m_dialect.empty() ? nullptr : m_dialect.c_str());
            const bool bResultSet = poLayer != nullptr;
            poSrcDS->ReleaseResultSet(poLayer);
            if (bResultSet && !m_quiet)
            {
                ReportError(CE_Warning, CPLE_AppDefined,
                            "Execution of the SQL statement '%s' returned a "
                            "result set. It will be ignored. You may silence "
                            "this warning with the 'quiet' argument.",
                            sql.c_str());
            }
            else if (CPLGetErrorCounter() > nErrorCounter &&
                     CPLGetLastErrorType() == CE_Failure)
            {
                ReportError(CE_Failure, CPLE_AppDefined,
                            "Execution of the SQL statement '%s' failed.%s",
                            sql.c_str(),
                            m_update ? ""
                                     : " Perhaps you need to specify the "
                                       "'update' argument?");
                return false;
            }
        }
        return true;
    }

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

    if (!m_outputLayer.empty() && m_outputLayer.size() != m_sql.size())
    {
        ReportError(CE_Failure, CPLE_AppDefined,
                    "There should be as many layer names in --output-layer as "
                    "in --statement");
        return false;
    }

    if (m_sql.size() == 1)
    {
        auto outDS = std::make_unique<GDALVectorSQLAlgorithmDataset>(*poSrcDS);
        outDS->SetDescription(poSrcDS->GetDescription());

        const auto nErrorCounter = CPLGetErrorCounter();
        OGRLayer *poLayer = poSrcDS->ExecuteSQL(
            m_sql[0].c_str(), nullptr,
            m_dialect.empty() ? nullptr : m_dialect.c_str());
        if (!poLayer)
        {
            if (nErrorCounter == CPLGetErrorCounter())
            {
                ReportError(CE_Failure, CPLE_AppDefined,
                            "Execution of the SQL statement '%s' did not "
                            "result in a result layer.",
                            m_sql[0].c_str());
            }
            return false;
        }

        if (!m_outputLayer.empty())
        {
            const std::string &osLayerName = m_outputLayer[0];
            poLayer->GetLayerDefn()->SetName(osLayerName.c_str());
            poLayer->SetDescription(osLayerName.c_str());
        }
        outDS->AddLayer(poLayer);
        m_outputDataset.Set(std::move(outDS));
    }
    else
    {
        // First pass to check all statements are valid and figure out layer
        // names
        std::set<std::string> setOutputLayerNames;
        std::vector<std::string> aosLayerNames;
        for (const std::string &sql : m_sql)
        {
            const auto nErrorCounter = CPLGetErrorCounter();
            auto poLayer = poSrcDS->ExecuteSQL(
                sql.c_str(), nullptr,
                m_dialect.empty() ? nullptr : m_dialect.c_str());
            if (!poLayer)
            {
                if (nErrorCounter == CPLGetErrorCounter())
                {
                    ReportError(CE_Failure, CPLE_AppDefined,
                                "Execution of the SQL statement '%s' did not "
                                "result in a result layer.",
                                sql.c_str());
                }
                return false;
            }

            std::string osLayerName;

            if (!m_outputLayer.empty())
            {
                osLayerName = m_outputLayer[aosLayerNames.size()];
            }
            else
            {
                osLayerName = poLayer->GetDescription();
                for (int num = 2;
                     cpl::contains(setOutputLayerNames, osLayerName); ++num)
                {
                    osLayerName = poLayer->GetDescription();
                    osLayerName += std::to_string(num);
                }
            }

            if (!osLayerName.empty())
            {
                poLayer->GetLayerDefn()->SetName(osLayerName.c_str());
                poLayer->SetDescription(osLayerName.c_str());
            }
            setOutputLayerNames.insert(poLayer->GetDescription());
            aosLayerNames.push_back(poLayer->GetDescription());

            poSrcDS->ReleaseResultSet(poLayer);
        }

        auto outDS =
            std::make_unique<GDALVectorSQLAlgorithmDatasetMultiLayer>(*poSrcDS);
        outDS->SetDescription(poSrcDS->GetDescription());

        for (size_t i = 0; i < aosLayerNames.size(); ++i)
        {
            outDS->AddLayer(m_sql[i], m_dialect, aosLayerNames[i]);
        }

        m_outputDataset.Set(std::move(outDS));
    }

    return true;
}

GDALVectorSQLAlgorithmStandalone::~GDALVectorSQLAlgorithmStandalone() = default;

//! @endcond