/******************************************************************************
 *
 * 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 "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 <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}");
    if (osFormat != "pbf")
    {
        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 aosCenter(
        CSLTokenizeString2(oObj.GetString("center").c_str(), ",", 0));
    if (aosCenter.size() != 3)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "Expected 3 values for center");
        return false;
    }
    const double dfCenterLong = CPLAtof(aosCenter[0]);
    const double dfCenterLat = CPLAtof(aosCenter[1]);
    if (std::fabs(dfCenterLong) > 180 || std::fabs(dfCenterLat) > 90)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "Invalid center");
        return false;
    }
    const int nCenterZoom = atoi(aosCenter[2]);
    if (nCenterZoom < 0 || nCenterZoom > 255)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "Missing or invalid center zoom");
        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;
    }

    CPLJSONDocument oMetadataDoc;
    oMetadataDoc.SetRoot(oObj);
    osMetadata = oMetadataDoc.SaveAsString();
    // CPLDebugOnly("PMTiles", "Metadata = %s", osMetadata.c_str());

    sHeader.root_dir_offset = 127;
    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 = pmtiles::COMPRESSION_GZIP;
    sHeader.tile_type = pmtiles::TILETYPE_MVT;
    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;
}

/************************************************************************/
/*                               HashArray()                            */
/************************************************************************/

// From https://codereview.stackexchange.com/questions/171999/specializing-stdhash-for-stdarray
// We do not use std::hash<std::array<T, N>> as the name of the struct
// because with gcc 5.4 we get the following error:
// https://stackoverflow.com/questions/25594644/warning-specialization-of-template-in-different-namespace
template <class T, size_t N> struct HashArray
{
    CPL_NOSANITIZE_UNSIGNED_INT_OVERFLOW
    size_t operator()(const std::array<T, N> &key) const
    {
        std::hash<T> hasher;
        size_t result = 0;
        for (size_t i = 0; i < N; ++i)
        {
            result = result * 31 + hasher(key[i]);
        }
        return result;
    }
};

/************************************************************************/
/*                    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 build a temporary file that contains the 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::string osTmpFile(std::string(pszDestName) + ".tmp");
    if (!VSIIsLocal(pszDestName))
    {
        osTmpFile = CPLGenerateTempFilename(CPLGetFilename(pszDestName));
    }

    auto poTmpFile =
        VSIVirtualHandleUniquePtr(VSIFOpenL(osTmpFile.c_str(), "wb+"));
    VSIUnlink(osTmpFile.c_str());
    if (!poTmpFile)
    {
        CPLError(CE_Failure, CPLE_FileIO, "Cannot open %s for write",
                 osTmpFile.c_str());
        return false;
    }

    struct ResetAndUnlinkTmpFile
    {
        VSIVirtualHandleUniquePtr &m_poFile;
        std::string m_osFilename;

        ResetAndUnlinkTmpFile(VSIVirtualHandleUniquePtr &poFile,
                              const std::string &osFilename)
            : m_poFile(poFile), m_osFilename(osFilename)
        {
        }

        ~ResetAndUnlinkTmpFile()
        {
            m_poFile.reset();
            VSIUnlink(m_osFilename.c_str());
        }
    };

    ResetAndUnlinkTmpFile oReseer(poTmpFile, osTmpFile);

    std::vector<pmtiles::entryv3> asPMTilesEntries;
    uint64_t nLastTileId = 0;
    uint64_t nFileOffset = 0;
    std::array<unsigned char, 16> abyLastMD5;
    std::unordered_map<std::array<unsigned char, 16>,
                       std::pair<uint64_t, uint32_t>,
                       HashArray<unsigned char, 16>>
        oMapMD5ToOffsetLen;
    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;

                if (poTmpFile->Write(pabyData, nTileDataLength, 1) != 1)
                {
                    CPLError(CE_Failure, CPLE_FileIO, "Failed writing");
                    return false;
                }
            }

            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;

    // Nomber 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 final 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 (poTmpFile->Seek(0, SEEK_SET) != 0 ||
        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 content of the temporary file at end of the output file.
    std::string oCopyBuffer;
    oCopyBuffer.resize(1024 * 1024);
    const uint64_t nTotalSize = nFileOffset;
    nFileOffset = 0;
    while (nFileOffset < nTotalSize)
    {
        const size_t nToRead = static_cast<size_t>(
            std::min<uint64_t>(nTotalSize - nFileOffset, oCopyBuffer.size()));
        if (poTmpFile->Read(&oCopyBuffer[0], nToRead, 1) != 1 ||
            poFile->Write(&oCopyBuffer[0], nToRead, 1) != 1)
        {
            CPLError(CE_Failure, CPLE_FileIO, "Failed writing");
            return false;
        }
        nFileOffset += nToRead;
    }

    if (poFile->Close() != 0)
        return false;

    return true;
}