#include "pathfinding.h"

#include <raymath.h>

static i32 dst(Vec2i a, Vec2i b) {
    //i32 dX = a.x - b.x;
    //i32 dY = a.y - b.y;
    //return dX * dX + dY * dY;
    int dstX = a.x - b.x;
    int dstY = a.y - b.y;
    if (dstX < 0) dstX = -dstX;
    if (dstY < 0) dstY = -dstY;

    if (dstX > dstY)
        return 14 * dstY + 10 * (dstX - dstY);
    return 14 * dstX + 10 * (dstY - dstX);
}

static PathData *pushPathWaypoint(PathData *pathData, Position waypoint, BzObjectPool *pathPool) {
    if (pathData->numWaypoints + 1 > PATH_DATA_SIZE) {
        PathData *newPathData = bzObjectPool(pathPool);
        BZ_ASSERT(newPathData);
        bzMemSet(newPathData, 0, sizeof(*newPathData));
        newPathData->numWaypoints = 0;
        newPathData->next = pathData;
        pathData = newPathData;
    }
    pathData->waypoints[pathData->numWaypoints++] = waypoint;
    return pathData;
}

static void reversePath(PathData *pathData) {
    while (pathData) {
        for (i32 i = 0; i < pathData->numWaypoints / 2; i++) {
            i32 left = i;
            i32 right = (i32) (pathData->numWaypoints - 1 - i);

            Position tmp = pathData->waypoints[left];
            pathData->waypoints[left] = pathData->waypoints[right];
            pathData->waypoints[right] = tmp;
        }
        pathData = pathData->next;
    }
}

static void smoothPath(BzTileMap *map, PathData *pathData, BzObjectPool *pool) {
    // Our smoothed path
    PathData *outPath = pathData;
    Position outPos = outPath->waypoints[0];
    size_t outIdx = 1;
    Position lastPos = outPos;

#define NEXT_WAYPOINT(path, idx, len)                       \
do {                                                        \
    idx++;                                                  \
    if (idx >= len) {                                       \
        path = path->next;                                  \
        idx = 0;                                            \
        if (path) len = path->numWaypoints;                 \
    }                                                       \
} while (0)

    PathData *currPath = pathData;
    PathData *nextPath = pathData;
    size_t currIdx = 0;
    size_t nextIdx = 0;
    // Needed, because we overwrite numWaypoints
    size_t currPathLen = currPath->numWaypoints;
    size_t nextPathLen = nextPath->numWaypoints;

    // Second element
    NEXT_WAYPOINT(currPath, currIdx, currPathLen);
    // Third element
    NEXT_WAYPOINT(nextPath, nextIdx, nextPathLen);
    NEXT_WAYPOINT(nextPath, nextIdx, nextPathLen);

    outPath->numWaypoints = 1;
    while (nextPath && nextIdx < nextPathLen) {
        Position currPos = currPath->waypoints[currIdx];
        Position nextPos = nextPath->waypoints[nextIdx];
        lastPos = nextPos;

        NEXT_WAYPOINT(currPath, currIdx, currPathLen);
        NEXT_WAYPOINT(nextPath, nextIdx, nextPathLen);

        if (!bzTileMapCanRayCastLine(map, outPos, nextPos)) {
            outPos = currPos;
            outPath->waypoints[outIdx++] = currPos;
            outPath->numWaypoints = outIdx;
            if (outIdx >= PATH_DATA_SIZE) {
                outPath = outPath->next;
                outIdx = 0;
                outPath->numWaypoints = 0;
            }

        }
    }
#undef NEXT_WAYPOINT
    outPath->waypoints[outIdx++] = lastPos;
    outPath->numWaypoints = outIdx;

    // Free old path
    pathData = outPath->next;
    while (pathData) {
        bzObjectPoolRelease(pool, pathData);
        pathData = pathData->next;
    }
    outPath->next = NULL;

}

typedef struct Heap {
    PathNode *arr;
    i32 size;
    i32 capacity;
    i32 mapWidth; // For calculating record idx
    PathNodeRecord *records;
} Heap;

static void heapPush(Heap *heap, PathNode node);
static PathNode heapPop(Heap *heap);
static void heapUpdate(Heap *heap, i32 idx);

bool pathfindAStar(const PathfindingDesc *desc) {
    BZ_ASSERT(desc->map);
    BZ_ASSERT(desc->alloc);
    BzTileMap *map = desc->map;

    Vec2i start = bzTileMapPosToTile(map, desc->start);
    Vec2i target = bzTileMapPosToTile(map, desc->target);
    BZ_ASSERT(start.x >= 0 && start.x < map->width);
    BZ_ASSERT(start.y >= 0 && start.y < map->height);

    // Perform very cheap ray cast check
    if (bzTileMapCanRayCastLine(map, desc->start, desc->target)) {
        PathData *pathData = bzObjectPool(desc->pool);
        BZ_ASSERT(pathData);
        pathData->waypoints[0] = desc->target;
        pathData->numWaypoints = 1;
        pathData->next         = NULL;

        if (desc->outPath)
            *desc->outPath = (Path) {pathData, 0};
        return true;
    }

    i32 numTiles = map->width * map->height;

    PathNodeRecord *closedSet = bzStackAlloc(desc->alloc, sizeof(*closedSet) * numTiles);
    bzMemSet(closedSet, 0, sizeof(*closedSet) * numTiles);

    PathNode *openSetArr = bzStackAlloc(desc->alloc, sizeof(*openSetArr) * numTiles);
    Heap openSet = {
            .arr=openSetArr,
            .size=0,
            .capacity=numTiles,
            .mapWidth=map->width,
            .records=closedSet
    };

    i32 toTargetCost = dst(start, target);
    heapPush(&openSet, (PathNode) {
            .weight=toTargetCost,
            .gCost=0,
            .hCost=toTargetCost,
            .pos=start
    });
    closedSet[start.y * map->width + start.x].open = true;

    bool foundPath = false;

    while (openSet.size) {
        PathNode node = heapPop(&openSet);
        if (node.pos.x == target.x &&
            node.pos.y == target.y) {
            // Found path
            foundPath = true;
            break;
        }
        Vec2i pos = node.pos;
        PathNodeRecord *nodeRecord = &closedSet[pos.y * map->width + pos.x];
        nodeRecord->visited = true;
        nodeRecord->open = false;


        // Node edges
        for (int y = node.pos.y - 1; y <= node.pos.y + 1; y++) {
            for (int x = node.pos.x - 1; x <= node.pos.x + 1; x++) {
                if (x == node.pos.x && y == node.pos.y)
                    continue;
                if (y < 0 || y >= map->height ||
                    x < 0 || x >= map->width)
                    continue;
                // not walkable
                if (bzTileMapHasAnyCollision(map, x, y) && x != target.x && y != target.y)
                    continue;
                PathNodeRecord *curRecord = &closedSet[y * map->width + x];
                if (curRecord->visited)
                    continue;

                Vec2i curPos = {x, y};
                i32 gCost = node.gCost + dst(node.pos, curPos);
                if (curRecord->open) {
                    i32 nodeIdx = curRecord->nodeIdx;
                    PathNode *curNode = &openSet.arr[nodeIdx];
                    if (gCost < curNode->gCost) {
                        curNode->gCost = gCost;
                        curNode->weight = curNode->gCost + curNode->hCost;
                        curRecord->toParentX = (i8) (curPos.x - node.pos.x);
                        curRecord->toParentY = (i8) (curPos.y - node.pos.y);
                        BZ_ASSERT(curPos.x == curNode->pos.x && curPos.y == curNode->pos.y);
                        heapUpdate(&openSet, nodeIdx);
                    }
                }

                if (!curRecord->open) {
                    toTargetCost = dst(curPos, target);

                    curRecord->toParentX = (i8) (curPos.x - node.pos.x);
                    curRecord->toParentY = (i8) (curPos.y - node.pos.y);
                    curRecord->open = true;

                    heapPush(&openSet, (PathNode) {
                            .weight = gCost + toTargetCost,
                            .gCost = gCost,
                            .hCost = toTargetCost,
                            .pos = curPos
                    });
                }
            }
        }
    }

    i32 pathLen = 0;
    if (foundPath && desc->outPath) {
        Vec2i pos = target;
        Path *out = desc->outPath;
        out->curWaypoint = 0;
        BZ_ASSERT(desc->pool);
        PathData *pathData = bzObjectPool(desc->pool);
        BZ_ASSERT(pathData);
        bzMemSet(pathData, 0, sizeof(*pathData));

        pathData = pushPathWaypoint(pathData, desc->target, desc->pool);

        // Write path
        while (pos.x != start.x || pos.y != start.y) {
            Position waypoint = {
                    pos.x * map->tileWidth + map->tileWidth * 0.5f,
                    pos.y * map->tileHeight + map->tileHeight * 0.5f
            };

            if (pathLen != 0)
                pathData = pushPathWaypoint(pathData, waypoint, desc->pool);

            PathNodeRecord visit = closedSet[pos.y * map->width + pos.x];
            BZ_ASSERT(visit.toParentX != 0 || visit.toParentY != 0);
            pos.x -= visit.toParentX;
            pos.y -= visit.toParentY;
            pathLen++;
        }
        if (pathLen == 0) {
            bzObjectPoolRelease(desc->pool, pathData);
            pathData = NULL;
        }

        reversePath(pathData);
        if (pathLen > 2)
            smoothPath(map, pathData, desc->pool);
        *desc->outPath = (Path) {pathData, 0};
    }

    bzStackAllocFree(desc->alloc, openSet.arr);
    bzStackAllocFree(desc->alloc, closedSet);

    return foundPath;
}

static void heapSwap(Heap *heap, i32 left, i32 right) {
    PathNode tmp = heap->arr[left];
    heap->arr[left] = heap->arr[right];
    heap->arr[right] = tmp;

    Vec2i leftPos = heap->arr[left].pos;
    Vec2i rightPos = heap->arr[right].pos;

    i32 leftIdx = leftPos.y * heap->mapWidth + leftPos.x;
    i32 rightIdx = rightPos.y * heap->mapWidth + rightPos.x;

    PathNodeRecord *leftNode = &heap->records[leftIdx];
    PathNodeRecord *rightNode = &heap->records[rightIdx];
    i32 tmpIdx = leftNode->nodeIdx;
    leftNode->nodeIdx = rightNode->nodeIdx;
    rightNode->nodeIdx = tmpIdx;
}

static i32 heapCmp(Heap *heap, i32 leftIdx, i32 rightIdx) {
    PathNode left = heap->arr[leftIdx];
    PathNode right = heap->arr[rightIdx];
    return left.weight - right.weight;
}

#define HEAP_LEFT(i) ((i32)(i * 2 + 1))
#define HEAP_RIGHT(i) ((i32)(i * 2 + 2))
#define HEAP_PARENT(i) ((i32) ((i - 1) / 2))

static void heapSiftUp(Heap *heap, i32 idx) {
    while (idx >= 0) {
        i32 parent = HEAP_PARENT(idx);
        if (heapCmp(heap, idx, parent) >= 0)
            break;
        heapSwap(heap, idx, parent);
        idx = parent;
    }

}
static void heapSiftDown(Heap *heap, i32 idx) {
    while (idx < heap->size) {
        i32 l = HEAP_LEFT(idx);
        i32 r = HEAP_RIGHT(idx);
        i32 smallest = idx;
        if (l < heap->size && heapCmp(heap, l, idx) < 0)
            smallest = l;
        if (r < heap->size && heapCmp(heap, r, smallest) < 0)
            smallest = r;

        if (smallest == idx) break;

        heapSwap(heap, idx, smallest);
        idx = smallest;

    }
}

static void heapPush(Heap *heap, PathNode node) {
    BZ_ASSERT(heap->size < heap->capacity);

    heap->arr[heap->size] = node;
    heap->size++;

    Vec2i pos = node.pos;
    PathNodeRecord *record = &heap->records[pos.y * heap->mapWidth + pos.x];
    record->nodeIdx = heap->size - 1;

    heapSiftUp(heap, heap->size - 1);
}
static PathNode heapPop(Heap *heap) {
    BZ_ASSERT(heap->size > 0);

    PathNode popped = heap->arr[0];
    heap->size--;

    heapSwap(heap, 0, heap->size);
    heapSiftDown(heap, 0);

    return popped;
}
static void heapUpdate(Heap *heap, i32 idx) {
    heapSiftUp(heap, idx);
    heapSiftDown(heap, idx);
}