#include "pathfinding.h"

static i32 dst(TilePosition a, TilePosition 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);
}

bool findPath(const PathfindingDesc *desc) {
    BZ_ASSERT(desc->map);
    BzTileMap *map = desc->map;
    BZ_ASSERT(desc->start.x >= 0 && desc->start.x < map->width);
    BZ_ASSERT(desc->start.y >= 0 && desc->start.y < map->height);
    typedef struct Visited {
        bool visited : 1;
        i8 x : 3;
        i8 y : 3;
    } Visited;
    Visited closedSet[map->width * map->height] = {};

    PathNode *openSet = desc->openSet;
    if (!openSet) openSet = bzHeapCreate(PathNode, map->width * map->height);
    else bzHeapClear(openSet);

    i32 toTargetCost = dst(desc->start, desc->target);
    bzHeapPush(openSet, (PathNode) {
        .weight=toTargetCost,
        .gCost=0,
        .hCost=toTargetCost,
        .visited=false,
        .pos=desc->start
    });

    bool foundPath = false;

    while (!bzHeapIsEmpty(openSet)) {
        PathNode node = bzHeapPop(openSet);
        if (node.pos.x == desc->target.x &&
            node.pos.y == desc->target.y) {
            // Found path
            foundPath = true;
            break;
        }

        // 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 (bzTileMapHasCollision(map, x, y))
                    continue;
                Visited *curClosed = &closedSet[y * map->width + x];
                if (curClosed->visited)
                    continue;
                curClosed->visited = true;

                TilePosition curPos = {x, y};
                i32 gCost = node.gCost + dst(node.pos, curPos);
                //if (gCost >= node.gCost) continue;
                toTargetCost = dst(curPos, desc->target);

                curClosed->x = (i8) (curPos.x - node.pos.x);
                curClosed->y = (i8) (curPos.y - node.pos.y);

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

            }
        }
    }

    i32 pathLen = 0;
    if (foundPath && desc->outPath) {
        TilePosition pos = desc->target;
        Path *out = desc->outPath;
        out->curWaypoint = 0;
        BZ_ASSERT(desc->pool);
        PathData *pathData = bzObjectPool(desc->pool);
        pathData->numWaypoints = 0;
        pathData->next = NULL;
        i32 numWaypoints = 0;
        // Write path
        // TODO: Write end pos
        while (pos.x != desc->start.x || pos.y != desc->start.y) {
            Position waypoint = {
                    pos.x * map->tileWidth + map->tileWidth * 0.5f,
                    pos.y * map->tileHeight + map->tileHeight * 0.5f
            };

            if (pathData->numWaypoints + 1 > PATH_DATA_SIZE) {
                PathData *newPathData = bzObjectPool(desc->pool);
                newPathData->numWaypoints = 0;
                newPathData->next = pathData;
                pathData = newPathData;
                numWaypoints = 0;
            }
            pathData->waypoints[numWaypoints++] = waypoint;
            pathData->numWaypoints = numWaypoints;

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

        // Reverse paths
        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;
        }
    }

    if (!desc->openSet)
        bzHeapDestroy(openSet);

    return foundPath ? pathLen : -1;
}