#include "pathfinding.h"

#include <raymath.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);
}

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 bool canRayCastLine(BzTileMap *map, Position from, Position to) {
    Vector2 step = Vector2Subtract(to, from);
    step = Vector2Normalize(step);

    while (Vector2DistanceSqr(from, to) > 1.0f) {
        from = Vector2Add(from, step);

        BzTile tileX = 0, tileY = 0;
        bzTileMapPosToTile(map, from, &tileX, &tileY);
        if (bzTileMapHasCollision(map, tileX, tileY))
            return false;
    }
    return true;
}

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;
    size_t outIdx = 1;
    Position outPos = outPath->waypoints[0];

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

    PathData *currPath = pathData;
    size_t currIdx = 0;
    PathData *nextPath = pathData;
    size_t nextIdx = 0;
    Position lastPos = outPos;
    // 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;

        currIdx++;
        nextIdx++;

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

        if (!canRayCastLine(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
    BZ_ASSERT(lastPos.x != INFINITY && lastPos.y != INFINITY);
    outPath->waypoints[outIdx++] = lastPos;
    outPath->numWaypoints = outIdx;

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

}

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

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

    PathNodeRecord *closedSet = desc->closedSet;
    if (!closedSet) closedSet = bzAlloc(sizeof(*closedSet) * map->width * map->height);
    bzMemSet(closedSet, 0, sizeof(*closedSet) * map->width * map->height);

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


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

    bool foundPath = false;

    while (!bzHeapIsEmpty(openSet)) {
        PathNode node = bzHeapPop(openSet);
        if (node.pos.x == target.x &&
            node.pos.y == target.y) {
            // Found path
            foundPath = true;
            break;
        }
        TilePosition 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 (bzTileMapHasCollision(map, x, y))
                    continue;
                PathNodeRecord *curRecord = &closedSet[y * map->width + x];
                if (curRecord->visited)
                    continue;

                TilePosition curPos = {x, y};
                i32 gCost = node.gCost + dst(node.pos, curPos);
#if 0
                if (curRecord->open) {
                    PathNode *curNode = NULL;
                    i32 curNodeIdx = -1;
                    for (i32 i = 0; i < bzHeapSize(openSet); i++) {
                        PathNode n = openSet[i];
                        if (curPos.x == n.pos.x && curPos.y == n.pos.y) {
                            curNode = &openSet[i];
                            curNodeIdx = i;
                            break;
                        }

                    }
                    BZ_ASSERT(curNode);
                    if (gCost < curNode->gCost) {
                        curNode->gCost = gCost;
                        curNode->weight = curNode->gCost + curNode->hCost;
                        curRecord->x = (i8) (curPos.x - node.pos.x);
                        curRecord->y = (i8) (curPos.y - node.pos.y);
                        BZ_ASSERT(curPos.x == curNode->pos.x && curPos.y == curNode->pos.y);
                        bzHeapUpdate(openSet, curNodeIdx);
                    }
                }
#endif

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

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

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

    i32 pathLen = 0;
    if (foundPath && desc->outPath) {
        TilePosition 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.x != 0 || visit.y != 0);
            pos.x -= visit.x;
            pos.y -= visit.y;
            pathLen++;
        }
        //pushPathWaypoint(pathData, desc->start, desc->pool);
        if (pathLen == 0) {
            bzObjectPoolRelease(desc->pool, pathData);
            pathData = NULL;
        }

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

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

    return foundPath ? pathLen : -1;
}