/*
 * Copyright 2020-2021 Uber Technologies, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
/** @file
 * @brief tests main H3 core library memory management.
 *
 *  usage: `testH3Memory`
 */

#include <math.h>
#include <string.h>

#include "cellsToMultiPoly.h"
#include "h3Index.h"
#include "h3api.h"
#include "latLng.h"
#include "linkedGeo.h"
#include "polyfill.h"
#include "polygon.h"
#include "test.h"
#include "utility.h"

// Whether to fail all allocations
static bool failAlloc = false;
// Actual number of malloc/calloc/realloc calls observed
static int actualAllocCalls = 0;
// Actual number of free calls observed
static int actualFreeCalls = 0;
// Set to non-zero to begin failing allocations after a certain number of calls
static int permittedAllocCalls = 0;

void resetMemoryCounters(int permitted) {
    failAlloc = false;
    actualAllocCalls = 0;
    actualFreeCalls = 0;
    permittedAllocCalls = permitted;
}

void *test_prefix_malloc(size_t size) {
    actualAllocCalls++;
    if (permittedAllocCalls && actualAllocCalls > permittedAllocCalls) {
        failAlloc = true;
    }
    if (failAlloc) {
        return NULL;
    }
    return malloc(size);
}

void *test_prefix_calloc(size_t num, size_t size) {
    actualAllocCalls++;
    if (permittedAllocCalls && actualAllocCalls > permittedAllocCalls) {
        failAlloc = true;
    }
    if (failAlloc) {
        return NULL;
    }
    return calloc(num, size);
}

void *test_prefix_realloc(void *ptr, size_t size) {
    actualAllocCalls++;
    if (permittedAllocCalls && actualAllocCalls > permittedAllocCalls) {
        failAlloc = true;
    }
    if (failAlloc) {
        return NULL;
    }
    return realloc(ptr, size);
}

void test_prefix_free(void *ptr) {
    actualFreeCalls++;
    free(ptr);
}

H3Index sunnyvale = 0x89283470c27ffff;
H3Index pentagon = 0x89080000003ffff;

static LatLng sfVerts[] = {
    {0.659966917655, -2.1364398519396},  {0.6595011102219, -2.1359434279405},
    {0.6583348114025, -2.1354884206045}, {0.6581220034068, -2.1382437718946},
    {0.6594479998527, -2.1384597563896}, {0.6599990002976, -2.1376771158464}};
static GeoLoop sfGeoLoop = {.numVerts = 6, .verts = sfVerts};
static GeoPolygon sfGeoPolygon;

SUITE(h3Memory) {
    TEST(gridDisk) {
        int k = 2;
        int64_t hexCount;
        t_assertSuccess(H3_EXPORT(maxGridDiskSize)(k, &hexCount));
        H3Index *gridDiskOutput = calloc(hexCount, sizeof(H3Index));

        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(gridDisk)(sunnyvale, k, gridDiskOutput));
        t_assert(actualAllocCalls == 0, "gridDisk did not call alloc");
        t_assert(actualFreeCalls == 0, "gridDisk did not call free");

        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(gridDisk)(pentagon, k, gridDiskOutput));
        t_assert(actualAllocCalls == 1, "gridDisk called alloc");
        t_assert(actualFreeCalls == 1, "gridDisk called free");

        resetMemoryCounters(0);
        failAlloc = true;
        memset(gridDiskOutput, 0, hexCount * sizeof(H3Index));
        t_assert(
            H3_EXPORT(gridDisk)(pentagon, k, gridDiskOutput) == E_MEMORY_ALLOC,
            "gridDisk returns E_MEMORY_ALLOC");
        t_assert(actualAllocCalls == 1, "gridDisk called alloc");
        t_assert(actualFreeCalls == 0, "gridDisk did not call free");

        for (int64_t i = 0; i < hexCount; i++) {
            t_assert(!gridDiskOutput[i],
                     "gridDisk did not produce output without alloc");
        }

        free(gridDiskOutput);
    }

    TEST(gridRing) {
        const int k = 2;
        int64_t ringSize;
        t_assertSuccess(H3_EXPORT(maxGridRingSize)(k, &ringSize));
        H3Index *gridRingOutput = calloc(ringSize, sizeof(H3Index));

        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(gridRing)(sunnyvale, k, gridRingOutput));
        t_assert(actualAllocCalls == 0, "gridRing did not call alloc");
        t_assert(actualFreeCalls == 0, "gridRing did not call free");

        resetMemoryCounters(2);
        t_assertSuccess(H3_EXPORT(gridRing)(pentagon, k, gridRingOutput));
        t_assert(actualAllocCalls == 2, "gridRing called alloc 2 times");
        t_assert(actualFreeCalls == 2, "gridRing called free 2 times");

        resetMemoryCounters(0);
        failAlloc = true;
        t_assert(
            H3_EXPORT(gridRing)(pentagon, k, gridRingOutput) == E_MEMORY_ALLOC,
            "gridRing returns E_MEMORY_ALLOC");
        t_assert(actualAllocCalls == 1, "gridRing called alloc 1 time");
        t_assert(actualFreeCalls == 0, "gridRing did not call free");

        resetMemoryCounters(1);
        t_assert(
            H3_EXPORT(gridRing)(pentagon, k, gridRingOutput) == E_MEMORY_ALLOC,
            "gridRing returns E_MEMORY_ALLOC");
        t_assert(actualAllocCalls == 2, "gridRing called alloc 2 times");
        t_assert(actualFreeCalls == 1, "gridRing called free 1 time");

        free(gridRingOutput);
    }

    TEST(compactCells) {
        int k = 9;
        int64_t hexCount;
        t_assertSuccess(H3_EXPORT(maxGridDiskSize)(k, &hexCount));
        int64_t expectedCompactCount = 73;

        // Generate a set of hexagons to compact
        H3Index *sunnyvaleExpanded = calloc(hexCount, sizeof(H3Index));
        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(gridDisk)(sunnyvale, k, sunnyvaleExpanded));
        t_assert(actualAllocCalls == 0, "gridDisk did not call alloc");
        t_assert(actualFreeCalls == 0, "gridDisk did not call free");

        H3Index *compressed = calloc(hexCount, sizeof(H3Index));

        resetMemoryCounters(0);
        failAlloc = true;
        H3Error err =
            H3_EXPORT(compactCells)(sunnyvaleExpanded, compressed, hexCount);
        t_assert(err == E_MEMORY_ALLOC, "malloc failed (1)");
        t_assert(actualAllocCalls == 1, "alloc called once");
        t_assert(actualFreeCalls == 0, "free not called");

        resetMemoryCounters(1);
        err = H3_EXPORT(compactCells)(sunnyvaleExpanded, compressed, hexCount);
        t_assert(err == E_MEMORY_ALLOC, "malloc failed (2)");
        t_assert(actualAllocCalls == 2, "alloc called twice");
        t_assert(actualFreeCalls == 1, "free called once");

        resetMemoryCounters(2);
        err = H3_EXPORT(compactCells)(sunnyvaleExpanded, compressed, hexCount);
        t_assert(err == E_MEMORY_ALLOC, "malloc failed (3)");
        t_assert(actualAllocCalls == 3, "alloc called three times");
        t_assert(actualFreeCalls == 2, "free called twice");

        resetMemoryCounters(3);
        err = H3_EXPORT(compactCells)(sunnyvaleExpanded, compressed, hexCount);
        t_assert(err == E_MEMORY_ALLOC, "compactCells failed (4)");
        t_assert(actualAllocCalls == 4, "alloc called four times");
        t_assert(actualFreeCalls == 3, "free called three times");

        resetMemoryCounters(4);
        err = H3_EXPORT(compactCells)(sunnyvaleExpanded, compressed, hexCount);
        t_assert(err == E_SUCCESS, "compact using successful malloc");
        t_assert(actualAllocCalls == 4, "alloc called four times");
        t_assert(actualFreeCalls == 4, "free called four times");

        int64_t count = 0;
        for (int64_t i = 0; i < hexCount; i++) {
            if (compressed[i] != 0) {
                count++;
            }
        }
        t_assert(count == expectedCompactCount, "got expected compact count");

        free(compressed);
        free(sunnyvaleExpanded);
    }

    TEST(polygonToCells) {
        sfGeoPolygon.geoloop = sfGeoLoop;
        sfGeoPolygon.numHoles = 0;

        int64_t numHexagons;
        t_assertSuccess(H3_EXPORT(maxPolygonToCellsSize)(&sfGeoPolygon, 9, 0,
                                                         &numHexagons));
        H3Index *hexagons = calloc(numHexagons, sizeof(H3Index));

        resetMemoryCounters(0);
        failAlloc = true;
        H3Error err = H3_EXPORT(polygonToCells)(&sfGeoPolygon, 9, 0, hexagons);
        t_assert(err == E_MEMORY_ALLOC, "polygonToCells failed (1)");
        t_assert(actualAllocCalls == 1, "alloc called once");
        t_assert(actualFreeCalls == 0, "free not called");

        resetMemoryCounters(1);
        err = H3_EXPORT(polygonToCells)(&sfGeoPolygon, 9, 0, hexagons);
        t_assert(err == E_MEMORY_ALLOC, "polygonToCells failed (2)");
        t_assert(actualAllocCalls == 2, "alloc called twice");
        t_assert(actualFreeCalls == 1, "free called once");

        resetMemoryCounters(2);
        err = H3_EXPORT(polygonToCells)(&sfGeoPolygon, 9, 0, hexagons);
        t_assert(err == E_MEMORY_ALLOC, "polygonToCells failed (3)");
        t_assert(actualAllocCalls == 3, "alloc called three times");
        t_assert(actualFreeCalls == 2, "free called twice");

        resetMemoryCounters(3);
        err = H3_EXPORT(polygonToCells)(&sfGeoPolygon, 9, 0, hexagons);
        t_assert(err == E_SUCCESS, "polygonToCells succeeded (4)");
        t_assert(actualAllocCalls == 3, "alloc called three times");
        t_assert(actualFreeCalls == 3, "free called three times");

        int64_t actualNumIndexes = countNonNullIndexes(hexagons, numHexagons);
        t_assert(actualNumIndexes == 1253, "got expected polygonToCells size");
        free(hexagons);
    }

    TEST(polygonToCellsExperimental) {
        sfGeoPolygon.geoloop = sfGeoLoop;
        sfGeoPolygon.numHoles = 0;

        int64_t numHexagons;
        t_assertSuccess(H3_EXPORT(maxPolygonToCellsSize)(
            &sfGeoPolygon, 9, CONTAINMENT_CENTER, &numHexagons));
        H3Index *hexagons = calloc(numHexagons, sizeof(H3Index));

        resetMemoryCounters(0);
        failAlloc = true;
        H3Error err = H3_EXPORT(polygonToCellsExperimental)(
            &sfGeoPolygon, 9, CONTAINMENT_CENTER, numHexagons, hexagons);
        t_assert(err == E_MEMORY_ALLOC,
                 "polygonToCellsExperimental failed (1)");
        t_assert(actualAllocCalls == 1, "alloc called once");
        t_assert(actualFreeCalls == 0, "free not called");

        resetMemoryCounters(1);
        err = H3_EXPORT(polygonToCellsExperimental)(
            &sfGeoPolygon, 9, CONTAINMENT_CENTER, numHexagons, hexagons);
        t_assert(err == E_SUCCESS, "polygonToCellsExperimental succeeded (1)");
        t_assert(actualAllocCalls == 1, "alloc called one time");
        t_assert(actualFreeCalls == 1, "free called one time");

        int64_t actualNumIndexes = countNonNullIndexes(hexagons, numHexagons);
        t_assert(actualNumIndexes == 1253,
                 "got expected polygonToCellsExperimental size");
        free(hexagons);
    }

    TEST(cellsToMultiPolygon) {
        // Exercise all error paths in cellsToMultiPolygon for coverage
        GeoMultiPolygon mpoly;
        H3Index cells2[] = {sunnyvale, 0x89283470c2fffff};
        H3Error err;

        // Find minimum allocations needed for success
        int minAllocs = 0;
        for (int permitted = 0; permitted < 50; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;  // Fail immediately
            err = H3_EXPORT(cellsToMultiPolygon)(cells2, 2, &mpoly);
            if (err == E_SUCCESS) {
                minAllocs = permitted;
                H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
                break;
            }
        }
        t_assert(minAllocs > 0, "Should require allocations");

        // Exercise error path at each allocation point for coverage
        for (int permitted = 0; permitted < minAllocs; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err = H3_EXPORT(cellsToMultiPolygon)(cells2, 2, &mpoly);
            t_assert(err == E_MEMORY_ALLOC, "Should fail before success");
        }
    }

    TEST(cellsToMultiPolygonWithHoles) {
        // Exercise error paths with polygons that have holes
        H3Index cells[] = {
            0x8027fffffffffff, 0x802bfffffffffff, 0x804dfffffffffff,
            0x8067fffffffffff, 0x806dfffffffffff, 0x8049fffffffffff,
        };
        int numCells = 6;
        GeoMultiPolygon mpoly;
        H3Error err;

        // Find minimum allocations and exercise all error paths
        for (int permitted = 0; permitted < 50; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err = H3_EXPORT(cellsToMultiPolygon)(cells, numCells, &mpoly);
            if (err == E_SUCCESS) {
                H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
                break;
            }
        }
    }

    TEST(cellsToMultiPolygonEmpty) {
        // Test the empty (globe) case
        H3Index *cells = NULL;
        GeoMultiPolygon mpoly;

        resetMemoryCounters(0);
        H3Error err = H3_EXPORT(cellsToMultiPolygon)(cells, 0, &mpoly);
        t_assert(err == E_SUCCESS, "Empty set should succeed");
        t_assert(mpoly.numPolygons == 0, "Empty set has no polygons");
        t_assert(mpoly.polygons == NULL, "Empty set has NULL polygons");
        t_assert(actualAllocCalls == 0, "Empty set needs no allocations");
    }

    TEST(cellsToMultiPolygonPartialFailure) {
        // Test failure after partial allocation to hit cleanup loops
        // Use a large diverse set of cells to create many loops and polygons
        H3Index cells[] = {
            // Multiple separate regions to create multiple polygons
            0x8027fffffffffff, 0x802bfffffffffff, 0x804dfffffffffff,
            0x8067fffffffffff, 0x806dfffffffffff, 0x8049fffffffffff,
            0x805ffffffffffff, 0x8057fffffffffff, 0x807dfffffffffff,
            0x80a5fffffffffff, 0x80a9fffffffffff, 0x808bfffffffffff,
            0x801bfffffffffff, 0x8035fffffffffff, 0x803ffffffffffff,
            0x8053fffffffffff, 0x8043fffffffffff, 0x8021fffffffffff,
            0x8011fffffffffff, 0x801ffffffffffff, 0x8097fffffffffff,
        };
        int numCells = 21;
        GeoMultiPolygon mpoly;
        H3Error err;

        // Find the allocation sequence and test every failure point
        int successPoint = 0;
        for (int permitted = 1; permitted < 150; permitted++) {
            resetMemoryCounters(permitted);
            err = H3_EXPORT(cellsToMultiPolygon)(cells, numCells, &mpoly);
            if (err == E_SUCCESS) {
                successPoint = permitted;
                H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
                break;
            }
        }

        // Test every single failure point to maximize branch coverage
        // This should hit all cleanup paths including:
        // - createSortableLoop failures after j > 0
        // - createMultiPolygon failures after createSortableLoopSet succeeds
        // - holes allocation failures after some polygons created
        for (int permitted = 1; permitted < successPoint; permitted++) {
            resetMemoryCounters(permitted);
            err = H3_EXPORT(cellsToMultiPolygon)(cells, numCells, &mpoly);
            t_assert(err == E_MEMORY_ALLOC, "Should fail with memory error");
        }
    }

    TEST(linkedGeoPolygonToGeoMultiPolygon) {
        // Generate a multi-polygon GeoMultiPolygon with holes from real
        // cells, then convert to linked form to use as test input
        H3Index cells[] = {
            0x8027fffffffffff, 0x802bfffffffffff, 0x804dfffffffffff,
            0x8067fffffffffff, 0x806dfffffffffff, 0x8049fffffffffff,
            0x8055fffffffffff,
        };
        GeoMultiPolygon src;
        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(cellsToMultiPolygon)(cells, 7, &src));

        LinkedGeoPolygon linked;
        resetMemoryCounters(0);
        t_assertSuccess(geoMultiPolygonToLinkedGeoPolygon(&src, &linked));
        H3_EXPORT(destroyGeoMultiPolygon)(&src);

        // Now test linkedGeoPolygonToGeoMultiPolygon failure paths
        GeoMultiPolygon mpoly;
        H3Error err;

        int minAllocs = 0;
        for (int permitted = 1; permitted < 50; permitted++) {
            resetMemoryCounters(permitted);
            err = linkedGeoPolygonToGeoMultiPolygon(&linked, &mpoly);
            if (err == E_SUCCESS) {
                minAllocs = permitted;
                H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
                break;
            }
        }
        t_assert(minAllocs > 0, "Should require allocations");

        for (int permitted = 0; permitted < minAllocs; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err = linkedGeoPolygonToGeoMultiPolygon(&linked, &mpoly);
            t_assert(err == E_MEMORY_ALLOC, "Should fail before success");
        }

        resetMemoryCounters(0);
        H3_EXPORT(destroyLinkedMultiPolygon)(&linked);
    }

    TEST(geoMultiPolygonToLinkedGeoPolygon) {
        // Generate a multi-polygon GeoMultiPolygon with holes from real cells
        H3Index cells[] = {
            0x8027fffffffffff, 0x802bfffffffffff, 0x804dfffffffffff,
            0x8067fffffffffff, 0x806dfffffffffff, 0x8049fffffffffff,
            0x8055fffffffffff,
        };
        GeoMultiPolygon mpoly;
        resetMemoryCounters(0);
        t_assertSuccess(H3_EXPORT(cellsToMultiPolygon)(cells, 7, &mpoly));

        LinkedGeoPolygon out;
        H3Error err;

        int minAllocs = 0;
        for (int permitted = 1; permitted < 50; permitted++) {
            resetMemoryCounters(permitted);
            err = geoMultiPolygonToLinkedGeoPolygon(&mpoly, &out);
            if (err == E_SUCCESS) {
                minAllocs = permitted;
                H3_EXPORT(destroyLinkedMultiPolygon)(&out);
                break;
            }
        }
        t_assert(minAllocs > 0, "Should require allocations");

        for (int permitted = 0; permitted < minAllocs; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err = geoMultiPolygonToLinkedGeoPolygon(&mpoly, &out);
            t_assert(err == E_MEMORY_ALLOC, "Should fail before success");
        }

        resetMemoryCounters(0);
        H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
    }

    TEST(cellsToLinkedMultiPolygon) {
        // Exercise allocation failure paths in cellsToLinkedMultiPolygon,
        // which composes cellsToMultiPolygon +
        // geoMultiPolygonToLinkedGeoPolygon. This catches double-free bugs when
        // the second phase fails and both the callee and caller attempt
        // cleanup.
        H3Index cells[] = {
            0x8027fffffffffff, 0x802bfffffffffff, 0x804dfffffffffff,
            0x8067fffffffffff, 0x806dfffffffffff, 0x8049fffffffffff,
            0x8055fffffffffff,
        };
        int numCells = ARRAY_SIZE(cells);
        LinkedGeoPolygon polygon;
        H3Error err;

        int minAllocs = 0;
        for (int permitted = 1; permitted < 100; permitted++) {
            resetMemoryCounters(permitted);
            err =
                H3_EXPORT(cellsToLinkedMultiPolygon)(cells, numCells, &polygon);
            if (err == E_SUCCESS) {
                minAllocs = permitted;
                H3_EXPORT(destroyLinkedMultiPolygon)(&polygon);
                break;
            }
        }
        t_assert(minAllocs > 0, "Should require allocations");

        for (int permitted = 0; permitted < minAllocs; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err =
                H3_EXPORT(cellsToLinkedMultiPolygon)(cells, numCells, &polygon);
            t_assert(err == E_MEMORY_ALLOC, "Should fail before success");
        }
    }

    TEST(cellsToMultiPolygonGlobe) {
        // Test with full set of resolution 0 cells
        // All edges cancel out, resulting in globe (8 triangular polygons)
        H3Index res0Cells[122];  // NUM_BASE_CELLS = 122
        GeoMultiPolygon mpoly;
        H3Error err;

        // Get all resolution 0 cells
        err = H3_EXPORT(getRes0Cells)(res0Cells);
        t_assert(err == E_SUCCESS, "getRes0Cells should succeed");

        // Find minimum allocations needed for success
        int minAllocs = 0;
        for (int permitted = 0; permitted < 100; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;  // Fail immediately
            err = H3_EXPORT(cellsToMultiPolygon)(res0Cells, 122, &mpoly);
            if (err == E_SUCCESS) {
                minAllocs = permitted;
                H3_EXPORT(destroyGeoMultiPolygon)(&mpoly);
                break;
            }
        }
        t_assert(minAllocs > 0, "Globe case should require allocations");

        // Exercise error path at each allocation point
        // This tests the createGlobeMultiPolygon error paths
        for (int permitted = 0; permitted < minAllocs; permitted++) {
            resetMemoryCounters(permitted);
            if (permitted == 0) failAlloc = true;
            err = H3_EXPORT(cellsToMultiPolygon)(res0Cells, 122, &mpoly);
            t_assert(err == E_MEMORY_ALLOC,
                     "Should fail with memory error before success");
        }
    }
}
