/******************************************************************************
 *
 * Project:  Elasticsearch Translator
 * Purpose:
 * Author:
 *
 ******************************************************************************
 * Copyright (c) 2021, Even Rouault <even dot rouault at spatialys.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 ****************************************************************************/

#include "ogr_elastic.h"
#include "ogrgeojsonreader.h"
#include "cpl_json.h"

#include <algorithm>
#include <set>

/************************************************************************/
/*                       OGRElasticAggregationLayer()                   */
/************************************************************************/

OGRElasticAggregationLayer::OGRElasticAggregationLayer(
    OGRElasticDataSource *poDS)
    : m_poDS(poDS)
{
    m_poFeatureDefn = new OGRFeatureDefn("aggregation");
    m_poFeatureDefn->SetGeomType(wkbPoint);
    m_poFeatureDefn->Reference();
    SetDescription(m_poFeatureDefn->GetName());

    OGRSpatialReference *poSRS_WGS84 = new OGRSpatialReference();
    poSRS_WGS84->SetFromUserInput(SRS_WKT_WGS84_LAT_LONG);
    poSRS_WGS84->SetAxisMappingStrategy(OAMS_TRADITIONAL_GIS_ORDER);
    m_poFeatureDefn->GetGeomFieldDefn(0)->SetSpatialRef(poSRS_WGS84);
    poSRS_WGS84->Dereference();

    OGRFieldDefn oKey("key", OFTString);
    m_poFeatureDefn->AddFieldDefn(&oKey);

    OGRFieldDefn oDocCount("doc_count", OFTInteger64);
    m_poFeatureDefn->AddFieldDefn(&oDocCount);
}

/************************************************************************/
/*                      ~OGRElasticAggregationLayer()                   */
/************************************************************************/

OGRElasticAggregationLayer::~OGRElasticAggregationLayer()
{
    m_poFeatureDefn->Release();
}

/************************************************************************/
/*                                  Build()                             */
/************************************************************************/

std::unique_ptr<OGRElasticAggregationLayer>
OGRElasticAggregationLayer::Build(OGRElasticDataSource *poDS,
                                  const char *pszAggregation)
{
    CPLJSONDocument oDoc;
    if (!oDoc.LoadMemory(pszAggregation))
        return nullptr;
    const auto oRoot = oDoc.GetRoot();
    const auto osIndex = oRoot.GetString("index");
    if (osIndex.empty())
    {
        CPLError(CE_Failure, CPLE_AppDefined,
                 "Missing 'index' member in AGGREGATION");
        return nullptr;
    }

    auto osGeometryField = oRoot.GetString("geometry_field");
    if (osGeometryField.empty())
    {
        std::set<CPLString> oSetLayers;
        std::vector<std::unique_ptr<OGRElasticLayer>> apoLayers;
        poDS->FetchMapping(osIndex.c_str(), oSetLayers, apoLayers);
        if (apoLayers.size() == 1)
        {
            apoLayers[0]->SetFeatureDefnFinalized();
            const int nGeomFieldCount =
                apoLayers[0]->GetLayerDefn()->GetGeomFieldCount();
            if (nGeomFieldCount == 1)
            {
                std::vector<CPLString> aosPath;
                bool bIsGeoPoint = false;
                apoLayers[0]->GetGeomFieldProperties(0, aosPath, bIsGeoPoint);
                for (const auto &osPart : aosPath)
                {
                    if (!osGeometryField.empty())
                        osGeometryField += '.';
                    osGeometryField += osPart;
                }
            }
            else if (nGeomFieldCount == 0)
            {
                CPLError(
                    CE_Failure, CPLE_AppDefined,
                    "No geometry field found upon which to build aggregation");
                return nullptr;
            }
            else
            {
                CPLError(CE_Failure, CPLE_AppDefined,
                         "Multiple geometry fields exist in the index. "
                         "Specify one with the 'geometry_field' member in "
                         "AGGREGATION");
                return nullptr;
            }
        }
        else
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Missing 'geometry_field' member in AGGREGATION");
            return nullptr;
        }
    }

    auto poLayer = std::unique_ptr<OGRElasticAggregationLayer>(
        new OGRElasticAggregationLayer(poDS));
    poLayer->m_osIndexName = osIndex;
    poLayer->m_osGeometryField = std::move(osGeometryField);

    // Parse geohash_grid options
    auto oGeohashGrid = oRoot["geohash_grid"];
    if (oGeohashGrid.IsValid() &&
        oGeohashGrid.GetType() == CPLJSONObject::Type::Object)
    {
        const int nPrecision = oGeohashGrid.GetInteger("precision");
        if (nPrecision > 0)
            poLayer->m_nGeohashGridPrecision = nPrecision;

        const int nMaxSize = oGeohashGrid.GetInteger("size");
        if (nMaxSize > 0)
            poLayer->m_nGeohashGridMaxSize = nMaxSize;
    }

    // Parse additional fields that correspond to statistical operations on
    // fields
    poLayer->m_oFieldDef = oRoot["fields"];
    if (poLayer->m_oFieldDef.IsValid() &&
        poLayer->m_oFieldDef.GetType() == CPLJSONObject::Type::Object)
    {
        // Start with stats, to keep track of the created columns, and
        // avoid duplicating them if a users ask for stats and min/max/etc.
        // on the same property.
        {
            auto oOp = poLayer->m_oFieldDef["stats"];
            if (oOp.IsValid() && oOp.GetType() == CPLJSONObject::Type::Array)
            {
                for (const auto &oField : oOp.ToArray())
                {
                    if (oField.GetType() == CPLJSONObject::Type::String)
                    {
                        for (const char *pszOp :
                             {"min", "max", "avg", "sum", "count"})
                        {
                            OGRFieldDefn oFieldDefn(
                                CPLSPrintf("%s_%s", oField.ToString().c_str(),
                                           pszOp),
                                strcmp(pszOp, "count") == 0 ? OFTInteger64
                                                            : OFTReal);
                            poLayer->m_poFeatureDefn->AddFieldDefn(&oFieldDefn);
                        }

                        CPLJSONObject oAgg;
                        CPLJSONObject oFieldAgg;
                        oFieldAgg.Add("field", oField.ToString());
                        oAgg.Add("stats", oFieldAgg);
                        poLayer->m_oAggregatedFieldsRequest.Add(
                            CPLSPrintf("%s_stats", oField.ToString().c_str()),
                            oAgg);
                    }
                }
            }
        }

        for (const char *pszOp : {"min", "max", "avg", "sum", "count"})
        {
            auto oOp = poLayer->m_oFieldDef[pszOp];
            if (oOp.IsValid() && oOp.GetType() == CPLJSONObject::Type::Array)
            {
                for (const auto &oField : oOp.ToArray())
                {
                    if (oField.GetType() == CPLJSONObject::Type::String)
                    {
                        const char *pszFieldName = CPLSPrintf(
                            "%s_%s", oField.ToString().c_str(), pszOp);
                        if (poLayer->m_poFeatureDefn->GetFieldIndex(
                                pszFieldName) < 0)
                        {
                            OGRFieldDefn oFieldDefn(pszFieldName,
                                                    strcmp(pszOp, "count") == 0
                                                        ? OFTInteger64
                                                        : OFTReal);
                            poLayer->m_poFeatureDefn->AddFieldDefn(&oFieldDefn);

                            CPLJSONObject oAgg;
                            CPLJSONObject oFieldAgg;
                            oFieldAgg.Add("field", oField.ToString());
                            oAgg.Add(strcmp(pszOp, "count") == 0 ? "value_count"
                                                                 : pszOp,
                                     oFieldAgg);
                            poLayer->m_oAggregatedFieldsRequest.Add(
                                oFieldDefn.GetNameRef(), oAgg);
                        }
                    }
                }
            }
        }
    }

    return poLayer;
}

/************************************************************************/
/*                          ResetReading()                              */
/************************************************************************/

void OGRElasticAggregationLayer::ResetReading()
{
    m_iCurFeature = 0;
}

/************************************************************************/
/*                          SetSpatialFilter()                          */
/************************************************************************/

void OGRElasticAggregationLayer::SetSpatialFilter(OGRGeometry *poGeom)

{
    OGRLayer::SetSpatialFilter(poGeom);
    m_bFeaturesRequested = false;
    m_apoCachedFeatures.clear();
}

/************************************************************************/
/*                          BuildRequest()                              */
/************************************************************************/

#define FILTERED_STR "filtered"
#define GRID_STR "grid"
#define CENTROID_STR "centroid"

// Return a JSON serialized document that is the payload to POST for a /_search
// request
std::string OGRElasticAggregationLayer::BuildRequest()
{
    CPLJSONDocument oDoc;
    auto oRoot = oDoc.GetRoot();
    oRoot.Set("size", 0);

    auto aggs = CPLJSONObject();

    // Clamp spatial filter if needed
    m_bRequestHasSpatialFilter = false;
    OGREnvelope sEnvelope;
    if (m_poFilterGeom)
    {
        m_poFilterGeom->getEnvelope(&sEnvelope);

        OGRElasticLayer::ClampEnvelope(sEnvelope);
        if (!(sEnvelope.MinX == -180 && sEnvelope.MinY == -90 &&
              sEnvelope.MaxX == 180 && sEnvelope.MaxY == 90))
        {
            m_bRequestHasSpatialFilter = true;
        }
    }

    if (m_bRequestHasSpatialFilter)
    {
        // Add spatial filtering
        auto top_aggs = CPLJSONObject();
        oRoot.Add("aggs", top_aggs);

        auto filtered = CPLJSONObject();
        top_aggs.Add(FILTERED_STR, filtered);

        auto filter = CPLJSONObject();
        filtered.Add("filter", filter);
        filtered.Add("aggs", aggs);

        auto geo_bounding_box = CPLJSONObject();
        filter.Add("geo_bounding_box", geo_bounding_box);

        auto coordinates = CPLJSONObject();
        geo_bounding_box.Add(m_osGeometryField, coordinates);

        auto top_left = CPLJSONObject();
        coordinates.Add("top_left", top_left);
        top_left.Add("lat", sEnvelope.MaxY);
        top_left.Add("lon", sEnvelope.MinX);

        auto bottom_right = CPLJSONObject();
        coordinates.Add("bottom_right", bottom_right);
        bottom_right.Add("lat", sEnvelope.MinY);
        bottom_right.Add("lon", sEnvelope.MaxX);
    }
    else
    {
        oRoot.Add("aggs", aggs);
    }

    auto grid = CPLJSONObject();
    aggs.Add(GRID_STR, grid);

    // Build geohash_grid aggregation object
    auto geohash_grid = CPLJSONObject();
    grid.Add("geohash_grid", geohash_grid);
    geohash_grid.Set("field", m_osGeometryField);

    if (m_nGeohashGridPrecision >= 1)
    {
        geohash_grid.Set("precision", m_nGeohashGridPrecision);
    }
    else if (!m_bRequestHasSpatialFilter || (sEnvelope.MinX < sEnvelope.MaxX &&
                                             sEnvelope.MinY < sEnvelope.MaxY))
    {
        const double dfSpatialRatio =
            m_bRequestHasSpatialFilter
                ? (sEnvelope.MaxX - sEnvelope.MinX) / 360. *
                      (sEnvelope.MaxY - sEnvelope.MinY) / 180.
                : 1.0;

        // A geohash of size 1 can encode up to 32 positions, size 2 up to 32*32
        // etc.
        const int geohashSize = static_cast<int>(std::min(
            12.0, std::max(1.0, log(m_nGeohashGridMaxSize / dfSpatialRatio) /
                                    log(32.0))));

        geohash_grid.Set("precision", geohashSize);
    }
    geohash_grid.Set("size", m_nGeohashGridMaxSize);

    auto subaggs = CPLJSONObject();
    grid.Add("aggs", subaggs);

    auto centroid = CPLJSONObject();
    subaggs.Add(CENTROID_STR, centroid);

    auto geo_centroid = CPLJSONObject();
    centroid.Add("geo_centroid", geo_centroid);

    geo_centroid.Set("field", m_osGeometryField);

    // Add extra fields
    for (const auto &oChild : m_oAggregatedFieldsRequest.GetChildren())
    {
        subaggs.Add(oChild.GetName(), oChild);
    }

    return oRoot.Format(CPLJSONObject::PrettyFormat::Plain);
}

/************************************************************************/
/*                       IssueAggregationRequest()                      */
/************************************************************************/

void OGRElasticAggregationLayer::IssueAggregationRequest()
{
    const auto IsNumericJsonType = [](json_type type)
    { return type == json_type_int || type == json_type_double; };

    m_apoCachedFeatures.clear();

    json_object *poResponse = m_poDS->RunRequest(
        (std::string(m_poDS->GetURL()) + "/" + m_osIndexName + "/_search")
            .c_str(),
        BuildRequest().c_str());
    if (!poResponse)
        return;
    json_object *poBuckets = json_ex_get_object_by_path(
        poResponse, m_bRequestHasSpatialFilter
                        ? "aggregations." FILTERED_STR "." GRID_STR ".buckets"
                        : "aggregations." GRID_STR ".buckets");
    if (poBuckets && json_object_get_type(poBuckets) == json_type_array)
    {
        const auto nBuckets = json_object_array_length(poBuckets);
        for (auto i = decltype(nBuckets){0}; i < nBuckets; i++)
        {
            json_object *poBucket = json_object_array_get_idx(poBuckets, i);
            if (poBucket && json_object_get_type(poBucket) == json_type_object)
            {
                OGRFeature *poFeature = new OGRFeature(m_poFeatureDefn);
                poFeature->SetFID(i);

                json_object *poKey =
                    CPL_json_object_object_get(poBucket, "key");
                if (poKey && json_object_get_type(poKey) == json_type_string)
                {
                    poFeature->SetField("key", json_object_get_string(poKey));
                }

                json_object *poDocCount =
                    CPL_json_object_object_get(poBucket, "doc_count");
                if (poDocCount &&
                    json_object_get_type(poDocCount) == json_type_int)
                {
                    poFeature->SetField("doc_count",
                                        static_cast<GIntBig>(
                                            json_object_get_int64(poDocCount)));
                }

                json_object *poLocation = json_ex_get_object_by_path(
                    poBucket, CENTROID_STR ".location");
                if (poLocation &&
                    json_object_get_type(poLocation) == json_type_object)
                {
                    json_object *poLat =
                        CPL_json_object_object_get(poLocation, "lat");
                    json_object *poLon =
                        CPL_json_object_object_get(poLocation, "lon");
                    if (poLat &&
                        IsNumericJsonType(json_object_get_type(poLat)) &&
                        poLon && IsNumericJsonType(json_object_get_type(poLon)))
                    {
                        OGRPoint *poPoint =
                            new OGRPoint(json_object_get_double(poLon),
                                         json_object_get_double(poLat));
                        poPoint->assignSpatialReference(
                            m_poFeatureDefn->GetGeomFieldDefn(0)
                                ->GetSpatialRef());
                        poFeature->SetGeometryDirectly(poPoint);
                    }
                }

                if (m_oFieldDef.IsValid() &&
                    m_oFieldDef.GetType() == CPLJSONObject::Type::Object)
                {
                    for (const char *pszOp :
                         {"min", "max", "avg", "sum", "count"})
                    {
                        auto oOp = m_oFieldDef[pszOp];
                        if (oOp.IsValid() &&
                            oOp.GetType() == CPLJSONObject::Type::Array)
                        {
                            for (const auto &oField : oOp.ToArray())
                            {
                                if (oField.GetType() ==
                                    CPLJSONObject::Type::String)
                                {
                                    json_object *poField =
                                        json_ex_get_object_by_path(
                                            poBucket,
                                            CPLSPrintf(
                                                "%s_%s.value",
                                                oField.ToString().c_str(),
                                                pszOp));
                                    if (poField &&
                                        IsNumericJsonType(
                                            json_object_get_type(poField)))
                                    {
                                        const char *pszFieldName = CPLSPrintf(
                                            "%s_%s", oField.ToString().c_str(),
                                            pszOp);
                                        if (strcmp(pszOp, "count") == 0)
                                        {
                                            poFeature->SetField(
                                                pszFieldName,
                                                static_cast<GIntBig>(
                                                    json_object_get_int64(
                                                        poField)));
                                        }
                                        else
                                        {
                                            poFeature->SetField(
                                                pszFieldName,
                                                json_object_get_double(
                                                    poField));
                                        }
                                    }
                                }
                            }
                        }
                    }

                    auto oOp = m_oFieldDef["stats"];
                    if (oOp.IsValid() &&
                        oOp.GetType() == CPLJSONObject::Type::Array)
                    {
                        for (const auto &oField : oOp.ToArray())
                        {
                            if (oField.GetType() == CPLJSONObject::Type::String)
                            {
                                for (const char *pszOp :
                                     {"min", "max", "avg", "sum", "count"})
                                {
                                    json_object *poField =
                                        json_ex_get_object_by_path(
                                            poBucket,
                                            CPLSPrintf(
                                                "%s_stats.%s",
                                                oField.ToString().c_str(),
                                                pszOp));
                                    if (poField &&
                                        IsNumericJsonType(
                                            json_object_get_type(poField)))
                                    {
                                        const char *pszFieldName = CPLSPrintf(
                                            "%s_%s", oField.ToString().c_str(),
                                            pszOp);
                                        if (strcmp(pszOp, "count") == 0)
                                        {
                                            poFeature->SetField(
                                                pszFieldName,
                                                static_cast<GIntBig>(
                                                    json_object_get_int64(
                                                        poField)));
                                        }
                                        else
                                        {
                                            poFeature->SetField(
                                                pszFieldName,
                                                json_object_get_double(
                                                    poField));
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                m_apoCachedFeatures.emplace_back(poFeature);
            }
        }
    }
    json_object_put(poResponse);
}

/************************************************************************/
/*                        GetNextRawFeature()                           */
/************************************************************************/

OGRFeature *OGRElasticAggregationLayer::GetNextRawFeature()
{
    if (!m_bFeaturesRequested)
    {
        m_bFeaturesRequested = true;
        IssueAggregationRequest();
    }
    if (m_iCurFeature < static_cast<int>(m_apoCachedFeatures.size()))
    {
        auto poFeature = m_apoCachedFeatures[m_iCurFeature]->Clone();
        ++m_iCurFeature;
        return poFeature;
    }

    return nullptr;
}

/************************************************************************/
/*                          GetFeatureCount()                           */
/************************************************************************/

GIntBig OGRElasticAggregationLayer::GetFeatureCount(int bForce)
{
    if (m_poFilterGeom == nullptr && m_poAttrQuery == nullptr)
    {
        if (!m_bFeaturesRequested)
        {
            m_bFeaturesRequested = true;
            IssueAggregationRequest();
        }
        return static_cast<int>(m_apoCachedFeatures.size());
    }
    return OGRLayer::GetFeatureCount(bForce);
}

/************************************************************************/
/*                          TestCapability()                            */
/************************************************************************/

int OGRElasticAggregationLayer::TestCapability(const char *pszCap)
{
    return EQUAL(pszCap, OLCStringsAsUTF8);
}

/************************************************************************/
/*                             GetDataset()                             */
/************************************************************************/

GDALDataset *OGRElasticAggregationLayer::GetDataset()
{
    return m_poDS;
}