#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);
        bzMemSet(newPathData, 0, sizeof(*newPathData));
        BZ_ASSERT(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];

    PathData *prevPath = pathData;
    size_t prevIdx = 1;
    PathData *nextPath = pathData;
    size_t nextIdx = 2;
    Position lastPos;
    // Needed, because we overwrite numWaypoints
    size_t currPathLen = prevPath->numWaypoints;
    size_t nextPathLen = nextPath->numWaypoints;
    outPath->numWaypoints = 1;
    while (nextPath && nextIdx < nextPathLen) {
        Position currPos = prevPath->waypoints[prevIdx];
        Position nextPos = nextPath->waypoints[nextIdx];
        lastPos = nextPos;

        prevIdx++;
        nextIdx++;

        if (prevIdx >= currPathLen) {
            prevPath = prevPath->next;
            prevIdx = 0;
            if (prevPath) currPathLen = prevPath->numWaypoints;
        }
        if (nextIdx >= nextPathLen) {
            nextPath = nextPath->next;
            nextIdx = 0;
            if (nextPath) nextPathLen = nextPath->numWaypoints;
        }
        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;
            }

        }
    }
    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 findPath(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);

    PathClosedNode *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,
        .visited=false,
        .pos=start
    });

    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;
        }

        // 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;
                PathClosedNode *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, 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 = target;
        Path *out = desc->outPath;
        out->curWaypoint = 0;
        BZ_ASSERT(desc->pool);
        PathData *pathData = bzObjectPool(desc->pool);
        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);

            PathClosedNode 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;
}