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

 *

 * Project:  OpenGIS Simple Features Reference Implementation

 * Purpose:  Implementation of PMTiles

 * Author:   Even Rouault <even.rouault at spatialys.com>

 *

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

 * Copyright (c) 2023, Planet Labs

 *

 * SPDX-License-Identifier: MIT

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

#include "cpl_json.h"

#include "ogrsf_frmts.h"

#include "ogr_pmtiles.h"

#include "ogrpmtilesfrommbtiles.h"

#include "include_pmtiles.h"

#include "cpl_compressor.h"

#include "cpl_md5.h"

#include "cpl_string.h"

#include "cpl_vsi_virtual.h"

#include <algorithm>

#include <array>

#include <cassert>

#include <unordered_map>

#include <unordered_set>

#include <utility>

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

/*                          ProcessMetadata()                           */

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

static bool ProcessMetadata(GDALDataset *poSQLiteDS, pmtiles::headerv3 &sHeader,

                            std::string &osMetadata)

{

    auto poMetadata = poSQLiteDS->GetLayerByName("metadata");

    if (!poMetadata)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "metadata table not found");

        return false;

    }

    const int iName = poMetadata->GetLayerDefn()->GetFieldIndex("name");

    const int iValue = poMetadata->GetLayerDefn()->GetFieldIndex("value");

    if (iName < 0 || iValue < 0)

    {

        CPLError(CE_Failure, CPLE_AppDefined,

                 "Bad structure for metadata table");

        return false;

    }

    CPLJSONObject oObj;

    CPLJSONDocument oJsonDoc;

    for (auto &&poFeature : poMetadata)

    {

        const char *pszName = poFeature->GetFieldAsString(iName);

        const char *pszValue = poFeature->GetFieldAsString(iValue);

        if (EQUAL(pszName, "json"))

        {

            if (!oJsonDoc.LoadMemory(pszValue))

            {

                CPLError(CE_Failure, CPLE_AppDefined,

                         "Cannot parse 'json' metadata item");

                return false;

            }

            for (const auto &oChild : oJsonDoc.GetRoot().GetChildren())

            {

                oObj.Add(oChild.GetName(), oChild);

            }

        }

        else

        {

            oObj.Add(pszName, pszValue);

        }

    }

    // MBTiles advertises scheme=tms. Override this

    oObj.Set("scheme", "xyz");

    const auto osFormat = oObj.GetString("format", "{missing}");

    uint8_t tile_type = pmtiles::TILETYPE_UNKNOWN;

    if (osFormat == "pbf")

        tile_type = pmtiles::TILETYPE_MVT;

    else if (osFormat == "png" || osFormat == "image/png")

        tile_type = pmtiles::TILETYPE_PNG;

    else if (osFormat == "jpg" || osFormat == "jpeg" ||

             osFormat == "image/jpeg")

        tile_type = pmtiles::TILETYPE_JPEG;

    else if (osFormat == "webp" || osFormat == "image/webp")

        tile_type = pmtiles::TILETYPE_WEBP;

    else if (osFormat == "avif" || osFormat == "image/avif")

        tile_type = pmtiles::TILETYPE_AVIF;

    else

    {

        CPLError(CE_Failure, CPLE_AppDefined, "format=%s unhandled",

                 osFormat.c_str());

        return false;

    }

    int nMinZoom = atoi(oObj.GetString("minzoom", "-1").c_str());

    if (nMinZoom < 0 || nMinZoom > 255)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Missing or invalid minzoom");

        return false;

    }

    int nMaxZoom = atoi(oObj.GetString("maxzoom", "-1").c_str());

    if (nMaxZoom < 0 || nMaxZoom > 255)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Missing or invalid maxzoom");

        return false;

    }

    const CPLStringList aosBounds(

        CSLTokenizeString2(oObj.GetString("bounds").c_str(), ",", 0));

    if (aosBounds.size() != 4)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Expected 4 values for bounds");

        return false;

    }

    const double dfMinX = CPLAtof(aosBounds[0]);

    const double dfMinY = CPLAtof(aosBounds[1]);

    const double dfMaxX = CPLAtof(aosBounds[2]);

    const double dfMaxY = CPLAtof(aosBounds[3]);

    if (std::fabs(dfMinX) > 180 || std::fabs(dfMinY) > 90 ||

        std::fabs(dfMaxX) > 180 || std::fabs(dfMaxY) > 90)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Invalid bounds");

        return false;

    }

    double dfCenterLong = 0;

    double dfCenterLat = 0;

    int nCenterZoom = 0;

    const auto osCenter = oObj.GetString("center");

    if (osCenter.empty())

    {

        dfCenterLong = (dfMinX + dfMaxX) / 2;

        dfCenterLat = (dfMinY + dfMaxY) / 2;

        nCenterZoom = nMaxZoom;

    }

    else

    {

        const CPLStringList aosCenter(

            CSLTokenizeString2(osCenter.c_str(), ",", 0));

        if (aosCenter.size() != 3)

        {

            CPLError(CE_Failure, CPLE_AppDefined,

                     "Expected 3 values for center");

            return false;

        }

        dfCenterLong = CPLAtof(aosCenter[0]);

        dfCenterLat = CPLAtof(aosCenter[1]);

        if (std::fabs(dfCenterLong) > 180 || std::fabs(dfCenterLat) > 90)

        {

            CPLError(CE_Failure, CPLE_AppDefined, "Invalid center");

            return false;

        }

        nCenterZoom = atoi(aosCenter[2]);

        if (nCenterZoom < 0 || nCenterZoom > 255)

        {

            CPLError(CE_Failure, CPLE_AppDefined,

                     "Missing or invalid center zoom");

            return false;

        }

    }

    CPLJSONDocument oMetadataDoc;

    oMetadataDoc.SetRoot(oObj);

    osMetadata = oMetadataDoc.SaveAsString();

    // CPLDebugOnly("PMTiles", "Metadata = %s", osMetadata.c_str());

    sHeader.root_dir_offset = PMTILES_HEADER_LENGTH;

    sHeader.root_dir_bytes = 0;

    sHeader.json_metadata_offset = 0;

    sHeader.json_metadata_bytes = 0;

    sHeader.leaf_dirs_offset = 0;

    sHeader.leaf_dirs_bytes = 0;

    sHeader.tile_data_offset = 0;

    sHeader.tile_data_bytes = 0;

    sHeader.addressed_tiles_count = 0;

    sHeader.tile_entries_count = 0;

    sHeader.tile_contents_count = 0;

    sHeader.clustered = true;

    sHeader.internal_compression = pmtiles::COMPRESSION_GZIP;

    sHeader.tile_compression = tile_type == pmtiles::TILETYPE_MVT

                                   ? pmtiles::COMPRESSION_GZIP

                                   : pmtiles::COMPRESSION_NONE;

    sHeader.tile_type = tile_type;

    sHeader.min_zoom = static_cast<uint8_t>(nMinZoom);

    sHeader.max_zoom = static_cast<uint8_t>(nMaxZoom);

    sHeader.min_lon_e7 = static_cast<int32_t>(dfMinX * 10e6);

    sHeader.min_lat_e7 = static_cast<int32_t>(dfMinY * 10e6);

    sHeader.max_lon_e7 = static_cast<int32_t>(dfMaxX * 10e6);

    sHeader.max_lat_e7 = static_cast<int32_t>(dfMaxY * 10e6);

    sHeader.center_zoom = static_cast<uint8_t>(nCenterZoom);

    sHeader.center_lon_e7 = static_cast<int32_t>(dfCenterLong * 10e6);

    sHeader.center_lat_e7 = static_cast<int32_t>(dfCenterLat * 10e6);

    return true;

}

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

/*                    OGRPMTilesConvertFromMBTiles()                    */

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

bool OGRPMTilesConvertFromMBTiles(const char *pszDestName,

                                  const char *pszSrcName)

{

    const char *const apszAllowedDrivers[] = {"SQLite", nullptr};

    auto poSQLiteDS = std::unique_ptr<GDALDataset>(

        GDALDataset::Open(pszSrcName, GDAL_OF_VECTOR, apszAllowedDrivers));

    if (!poSQLiteDS)

    {

        CPLError(CE_Failure, CPLE_AppDefined,

                 "Cannot open %s with SQLite driver", pszSrcName);

        return false;

    }

    pmtiles::headerv3 sHeader;

    std::string osMetadata;

    if (!ProcessMetadata(poSQLiteDS.get(), sHeader, osMetadata))

        return false;

    auto poTilesLayer = poSQLiteDS->GetLayerByName("tiles");

    if (!poTilesLayer)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "tiles table not found");

        return false;

    }

    const int iZoomLevel =

        poTilesLayer->GetLayerDefn()->GetFieldIndex("zoom_level");

    const int iTileColumn =

        poTilesLayer->GetLayerDefn()->GetFieldIndex("tile_column");

    const int iTileRow =

        poTilesLayer->GetLayerDefn()->GetFieldIndex("tile_row");

    const int iTileData =

        poTilesLayer->GetLayerDefn()->GetFieldIndex("tile_data");

    if (iZoomLevel < 0 || iTileColumn < 0 || iTileRow < 0 || iTileData < 0)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Bad structure for tiles table");

        return false;

    }

    struct TileEntry

    {

        uint64_t nTileId;

        std::array<unsigned char, 16> abyMD5;

    };

    // In a first step browse through the tiles table to compute the PMTiles

    // tile_id of each tile, and compute a hash of the tile data for

    // deduplication

    std::vector<TileEntry> asTileEntries;

    for (auto &&poFeature : poTilesLayer)

    {

        const int nZoomLevel = poFeature->GetFieldAsInteger(iZoomLevel);

        if (nZoomLevel < 0 || nZoomLevel > 30)

        {

            CPLError(CE_Warning, CPLE_AppDefined,

                     "Skipping tile with missing or invalid zoom_level");

            continue;

        }

        const int nColumn = poFeature->GetFieldAsInteger(iTileColumn);

        if (nColumn < 0 || nColumn >= (1 << nZoomLevel))

        {

            CPLError(CE_Warning, CPLE_AppDefined,

                     "Skipping tile with missing or invalid tile_column");

            continue;

        }

        const int nRow = poFeature->GetFieldAsInteger(iTileRow);

        if (nRow < 0 || nRow >= (1 << nZoomLevel))

        {

            CPLError(CE_Warning, CPLE_AppDefined,

                     "Skipping tile with missing or invalid tile_row");

            continue;

        }

        // MBTiles uses a 0=bottom-most row, whereas PMTiles uses

        // 0=top-most row

        const int nY = (1 << nZoomLevel) - 1 - nRow;

        uint64_t nTileId;

        try

        {

            nTileId = pmtiles::zxy_to_tileid(static_cast<uint8_t>(nZoomLevel),

                                             nColumn, nY);

        }

        catch (const std::exception &e)

        {

            // shouldn't happen given previous checks

            CPLError(CE_Failure, CPLE_AppDefined, "Cannot compute tile id: %s",

                     e.what());

            return false;

        }

        int nTileDataLength = 0;

        const GByte *pabyData =

            poFeature->GetFieldAsBinary(iTileData, &nTileDataLength);

        if (!pabyData)

        {

            CPLError(CE_Failure, CPLE_AppDefined, "Missing tile_data");

            return false;

        }

        TileEntry sEntry;

        sEntry.nTileId = nTileId;

        CPLMD5Context md5context;

        CPLMD5Init(&md5context);

        CPLMD5Update(&md5context, pabyData, nTileDataLength);

        CPLMD5Final(&sEntry.abyMD5[0], &md5context);

        try

        {

            asTileEntries.push_back(sEntry);

        }

        catch (const std::exception &e)

        {

            CPLError(CE_Failure, CPLE_AppDefined,

                     "Out of memory browsing through tiles: %s", e.what());

            return false;

        }

    }

    // Sort the tiles by ascending tile_id. This is a requirement to build

    // the PMTiles directories.

    std::sort(asTileEntries.begin(), asTileEntries.end(),

              [](const TileEntry &a, const TileEntry &b)

              { return a.nTileId < b.nTileId; });

    // Let's gather tile data in

    // a way that corresponds to the "clustered" mode, that is

    // "offsets are either contiguous with the previous offset+length, or

    // refer to a lesser offset, when writing with deduplication."

    std::vector<pmtiles::entryv3> asPMTilesEntries;

    uint64_t nFileOffset = 0;

    std::unordered_map<std::array<unsigned char, 16>,

                       std::pair<uint64_t, uint32_t>,

                       HashArray<unsigned char, 16>>

        oMapMD5ToOffsetLen;

    {

        uint64_t nLastTileId = 0;

        std::array<unsigned char, 16> abyLastMD5{0, 0, 0, 0, 0, 0, 0, 0,

                                                 0, 0, 0, 0, 0, 0, 0, 0};

        for (const auto &sEntry : asTileEntries)

        {

            if (sEntry.nTileId == nLastTileId + 1 &&

                sEntry.abyMD5 == abyLastMD5)

            {

                // If the tile id immediately follows the previous one and

                // has the same tile data, increase the run_length

                asPMTilesEntries.back().run_length++;

            }

            else

            {

                pmtiles::entryv3 sPMTilesEntry;

                sPMTilesEntry.tile_id = sEntry.nTileId;

                sPMTilesEntry.run_length = 1;

                auto oIter = oMapMD5ToOffsetLen.find(sEntry.abyMD5);

                if (oIter != oMapMD5ToOffsetLen.end())

                {

                    // Point to previously written tile data if this content

                    // has already been written

                    sPMTilesEntry.offset = oIter->second.first;

                    sPMTilesEntry.length = oIter->second.second;

                }

                else

                {

                    try

                    {

                        const auto sXYZ =

                            pmtiles::tileid_to_zxy(sEntry.nTileId);

                        poTilesLayer->SetAttributeFilter(CPLSPrintf(

                            "zoom_level = %d AND tile_column = %u AND tile_row "

                            "= "

                            "%u",

                            sXYZ.z, sXYZ.x, (1U << sXYZ.z) - 1U - sXYZ.y));

                    }

                    catch (const std::exception &e)

                    {

                        // shouldn't happen given previous checks

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Cannot compute xyz: %s", e.what());

                        return false;

                    }

                    poTilesLayer->ResetReading();

                    auto poFeature = std::unique_ptr<OGRFeature>(

                        poTilesLayer->GetNextFeature());

                    if (!poFeature)

                    {

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Cannot find tile");

                        return false;

                    }

                    int nTileDataLength = 0;

                    const GByte *pabyData = poFeature->GetFieldAsBinary(

                        iTileData, &nTileDataLength);

                    if (!pabyData)

                    {

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Missing tile_data");

                        return false;

                    }

                    sPMTilesEntry.offset = nFileOffset;

                    sPMTilesEntry.length = nTileDataLength;

                    oMapMD5ToOffsetLen[sEntry.abyMD5] =

                        std::pair<uint64_t, uint32_t>(nFileOffset,

                                                      nTileDataLength);

                    nFileOffset += nTileDataLength;

                }

                asPMTilesEntries.push_back(sPMTilesEntry);

                nLastTileId = sEntry.nTileId;

                abyLastMD5 = sEntry.abyMD5;

            }

        }

    }

    const CPLCompressor *psCompressor = CPLGetCompressor("gzip");

    assert(psCompressor);

    std::string osCompressed;

    struct compression_exception : std::exception

    {

        const char *what() const noexcept override

        {

            return "Compression failed";

        }

    };

    const auto oCompressFunc = [psCompressor,

                                &osCompressed](const std::string &osBytes,

                                               uint8_t) -> std::string

    {

        osCompressed.resize(32 + osBytes.size() * 2);

        size_t nOutputSize = osCompressed.size();

        void *pOutputData = &osCompressed[0];

        if (!psCompressor->pfnFunc(osBytes.data(), osBytes.size(), &pOutputData,

                                   &nOutputSize, nullptr,

                                   psCompressor->user_data))

        {

            throw compression_exception();

        }

        osCompressed.resize(nOutputSize);

        return osCompressed;

    };

    std::string osCompressedMetadata;

    std::string osRootBytes;

    std::string osLeaveBytes;

    int nNumLeaves;

    try

    {

        osCompressedMetadata =

            oCompressFunc(osMetadata, pmtiles::COMPRESSION_GZIP);

        // Build the root and leave directories (one depth max)

        std::tie(osRootBytes, osLeaveBytes, nNumLeaves) =

            pmtiles::make_root_leaves(oCompressFunc, pmtiles::COMPRESSION_GZIP,

                                      asPMTilesEntries);

    }

    catch (const std::exception &e)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Cannot build directories: %s",

                 e.what());

        return false;

    }

    // Finalize the header fields related to offsets and size of the

    // different parts of the file

    sHeader.root_dir_bytes = osRootBytes.size();

    sHeader.json_metadata_offset =

        sHeader.root_dir_offset + sHeader.root_dir_bytes;

    sHeader.json_metadata_bytes = osCompressedMetadata.size();

    sHeader.leaf_dirs_offset =

        sHeader.json_metadata_offset + sHeader.json_metadata_bytes;

    sHeader.leaf_dirs_bytes = osLeaveBytes.size();

    sHeader.tile_data_offset =

        sHeader.leaf_dirs_offset + sHeader.leaf_dirs_bytes;

    sHeader.tile_data_bytes = nFileOffset;

    // Number of tiles that are addressable in the PMTiles archive, that is

    // the number of tiles we would have if not deduplicating them

    sHeader.addressed_tiles_count = asTileEntries.size();

    // Number of tile entries in root and leave directories

    // ie entries whose run_length >= 1

    sHeader.tile_entries_count = asPMTilesEntries.size();

    // Number of distinct tile blobs

    sHeader.tile_contents_count = oMapMD5ToOffsetLen.size();

    // Now build the file!

    auto poFile = VSIVirtualHandleUniquePtr(VSIFOpenL(pszDestName, "wb"));

    if (!poFile)

    {

        CPLError(CE_Failure, CPLE_FileIO, "Cannot open %s for write",

                 pszDestName);

        return false;

    }

    const auto osHeader = sHeader.serialize();

    if (poFile->Write(osHeader.data(), osHeader.size(), 1) != 1 ||

        poFile->Write(osRootBytes.data(), osRootBytes.size(), 1) != 1 ||

        poFile->Write(osCompressedMetadata.data(), osCompressedMetadata.size(),

                      1) != 1 ||

        (!osLeaveBytes.empty() &&

         poFile->Write(osLeaveBytes.data(), osLeaveBytes.size(), 1) != 1))

    {

        CPLError(CE_Failure, CPLE_FileIO, "Failed writing");

        return false;

    }

    // Copy tile content at end of the output file.

    {

        uint64_t nLastTileId = 0;

        uint64_t nFileOffset2 = 0;

        std::array<unsigned char, 16> abyLastMD5{0, 0, 0, 0, 0, 0, 0, 0,

                                                 0, 0, 0, 0, 0, 0, 0, 0};

        std::unordered_set<std::array<unsigned char, 16>,

                           HashArray<unsigned char, 16>>

            oSetMD5;

        for (const auto &sEntry : asTileEntries)

        {

            if (sEntry.nTileId == nLastTileId + 1 &&

                sEntry.abyMD5 == abyLastMD5)

            {

                // If the tile id immediately follows the previous one and

                // has the same tile data, do nothing

            }

            else

            {

                auto oIter = oSetMD5.find(sEntry.abyMD5);

                if (oIter == oSetMD5.end())

                {

                    try

                    {

                        const auto sXYZ =

                            pmtiles::tileid_to_zxy(sEntry.nTileId);

                        poTilesLayer->SetAttributeFilter(CPLSPrintf(

                            "zoom_level = %d AND tile_column = %u AND tile_row "

                            "= "

                            "%u",

                            sXYZ.z, sXYZ.x, (1U << sXYZ.z) - 1U - sXYZ.y));

                    }

                    catch (const std::exception &e)

                    {

                        // shouldn't happen given previous checks

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Cannot compute xyz: %s", e.what());

                        return false;

                    }

                    poTilesLayer->ResetReading();

                    auto poFeature = std::unique_ptr<OGRFeature>(

                        poTilesLayer->GetNextFeature());

                    if (!poFeature)

                    {

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Cannot find tile");

                        return false;

                    }

                    int nTileDataLength = 0;

                    const GByte *pabyData = poFeature->GetFieldAsBinary(

                        iTileData, &nTileDataLength);

                    if (!pabyData)

                    {

                        CPLError(CE_Failure, CPLE_AppDefined,

                                 "Missing tile_data");

                        return false;

                    }

                    oSetMD5.insert(sEntry.abyMD5);

                    if (poFile->Write(pabyData, nTileDataLength, 1) != 1)

                    {

                        CPLError(CE_Failure, CPLE_FileIO, "Failed writing");

                        return false;

                    }

                    nFileOffset2 += nTileDataLength;

                }

                nLastTileId = sEntry.nTileId;

                abyLastMD5 = sEntry.abyMD5;

            }

        }

        CPL_IGNORE_RET_VAL(nFileOffset2);

        CPLAssert(nFileOffset2 == nFileOffset);

    }

    if (poFile->Close() != 0)

        return false;

    return true;

}