#include "systems.h"

#include "../game_state.h"
#include "../input.h"
#include "../buildings.h"
#include "../pathfinding.h"
#include "../unit_ai.h"
#include "../unit_actions.h"

#include <rlImGui.h>
#include <raymath.h>
#include <rlgl.h>

Rectangle calculateEntityBounds(Position pos, Size size);
bool getEntityBounds(ecs_entity_t entity, Position *outPos, Size *outSize, Rectangle *outBounds);

ecs_entity_t queryEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag);

bool pickEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag);
void pickUnits(BzSpatialGrid *entityGrid, Rectangle area);

static void iterateSelectedUnits(ecs_query_t *query, void (*fn)(ecs_entity_t entity, Position *pos, Size *size));
static void iterRemovePaths(ecs_entity_t entity, Position *pos, Size *size);

void placeUnits(i32 numUnits, f32 unitSpacing, Vector2 start, Vector2 end, BzTileMap *map, Vector2 **outPlaces);

void resetInputState(InputState *input) {
    input->cursor   = CURSOR_NONE;
    input->state    = INPUT_NONE;
    input->building = 0;
}

void inputPrimaryAction(Game *game, InputState *input) {
    const MouseButton primaryBtn = input->mapping.primaryBtn;
    if (isInputBtnDragged(input, primaryBtn)) {
        Vector2 start = input->mouseDownWorld;
        Vector2 end = input->mouseWorld;
        if (start.x > end.x) {
            f32 tmp = start.x;
            start.x = end.x;
            end.x = tmp;
        }
        if (start.y > end.y) {
            f32 tmp = start.y;
            start.y = end.y;
            end.y = tmp;
        }
        input->pickArea = (Rectangle) {start.x, start.y, end.x - start.x, end.y - start.y};
        pickUnits(game->entityGrid, input->pickArea);
    }
    i32 selectedCount = ecs_query_entity_count(input->queries.selected);
    if (isInputBtnJustDragged(input, primaryBtn) && selectedCount > 0) {
        resetInputState(input);
        input->state = INPUT_SELECTED_UNITS;
    } else if (isInputBtnJustUp(input, primaryBtn)) {
        if (pickEntity(game->entityGrid, input->mouseDownWorld, ecs_id(Unit))) {
            resetInputState(input);
            input->state = INPUT_SELECTED_UNITS;
        } else if (pickEntity(game->entityGrid, input->mouseDownWorld, Harvestable)) {
            resetInputState(input);
            input->state = INPUT_SELECTED_OBJECT;
        } else if (pickEntity(game->entityGrid, input->mouseDownWorld, Buildable)) {
            resetInputState(input);
            input->state = INPUT_SELECTED_BUILDING;
        }
        selectedCount = ecs_query_entity_count(input->queries.selected);
    }

    if (selectedCount == 0)
        resetInputState(input);
}
void inputUnitAction(Game *game, InputState *input) {
    ecs_query_t *query = input->queries.selected;
    BzTileMap   *map   = &game->map;
    const i32 numUnits = ecs_query_entity_count(query);
    BZ_ASSERT(numUnits > 0);

    const MouseButton actionBtn = input->mapping.secondaryBtn;

    input->cursor = CURSOR_NONE;
    ecs_entity_t taskEntity;
    if ((taskEntity = queryEntity(game->entityGrid, input->mouseWorld, Harvestable))) {
        input->cursor = CURSOR_COLLECT_WOOD;
        if (isInputBtnJustUp(input, actionBtn)) {
            ecs_defer_begin(ECS);
            ecs_iter_t it = ecs_query_iter(ECS, query);
            while (ecs_query_next(&it)) {
                for (i32 i = 0; i < it.count; i++) {
                    const ecs_entity_t entity = it.entities[i];
                    const Position target = *ecs_get(ECS, taskEntity, Position);
                    setUnitAI(entity, game, &(const UnitAI) {
                        .type = AI_WORKER_HARVEST,
                        .as.workerHarvest.resource = RES_WOOD,
                        .as.workerHarvest.target = taskEntity,
                        .as.workerHarvest.targetPosition = target
                    });
                    //addAction(entity, game, &(const Action) {
                    //    .type = ACTION_MOVE_TO,
                    //    .as.moveTo.target = target,
                    //    .as.moveTo.proximityThreshold = 10.0f,
                    //});
                    //ecs_set(ECS, entity, HarvestTask, {taskEntity});
                    goto while_break;
                }
            }
            while_break:
            ecs_iter_fini(&it);
            ecs_defer_end(ECS);
        }
        return;
    }

    if (isInputBtnJustUp(input, actionBtn)) {
        // Note: We mustn't use ecs ecs_remove_all since this will also
        // remove ongoing paths that are not part of this query.
        ecs_defer_begin(ECS);
        iterateSelectedUnits(query, iterRemovePaths);
        ecs_defer_end(ECS);

        const Position target = input->mouseWorld;

        ecs_iter_t it = ecs_query_iter(ECS, query);
        ecs_defer_begin(ECS);
        while (ecs_iter_next(&it)) {
            for (i32 i = 0; i < it.count; i++) {
                const ecs_entity_t entity = it.entities[i];
                clearActions(entity, game);
                addAction(entity, game, &(const Action) {
                    .type = ACTION_MOVE_TO,
                    .as.moveTo.target = target,
                    .as.moveTo.proximityThreshold = 6.0f,
                });
            }
        }
        ecs_defer_end(ECS);
    }

}

void updatePlayerInput() {
    f32 dt = GetFrameTime();
    ImGuiIO *io = igGetIO();
    if (io->WantCaptureMouse || io->WantCaptureKeyboard)
        return;

    Game *game = ecs_singleton_get_mut(ECS, Game);
    InputState *input = ecs_singleton_get_mut(ECS, InputState);

    const f32 maxZoom = 4.5f;
    const f32 minZoom = 0.9f;
    float moveSpeed = 100.0f * dt * (1 + maxZoom - game->camera.zoom);
    if (IsKeyDown(KEY_W)) game->camera.target.y -= moveSpeed;
    if (IsKeyDown(KEY_S)) game->camera.target.y += moveSpeed;
    if (IsKeyDown(KEY_A)) game->camera.target.x -= moveSpeed;
    if (IsKeyDown(KEY_D)) game->camera.target.x += moveSpeed;

    if (IsKeyDown(KEY_Q)) game->camera.rotation--;
    if (IsKeyDown(KEY_E)) game->camera.rotation++;

    // https://gamedev.stackexchange.com/questions/9330/zoom-to-cursor-calculation
    float wheel = GetMouseWheelMove();
    if (wheel != 0.0f) {
        const float zoomIncrement = 0.125f;
        f32 oldZoom = game->camera.zoom;
        f32 newZoom = oldZoom + wheel * zoomIncrement;
        newZoom = Clamp(newZoom, minZoom, maxZoom);

        Vector2 mouse = GetMousePosition();
        f32 mapWidth = GetScreenWidth() / oldZoom;
        f32 mapHeight = GetScreenHeight() / oldZoom;
        f32 widthRatio = (mouse.x - (GetScreenWidth() / 2)) / GetScreenWidth();

        f32 newMapWidth = GetScreenWidth() / newZoom;
        f32 newMapHeight = GetScreenHeight() / newZoom;
        f32 heightRatio = (mouse.y - (GetScreenHeight() / 2)) / GetScreenHeight();

        game->camera.target.x += (mapWidth - newMapWidth) * widthRatio;
        game->camera.target.y += (mapHeight - newMapHeight) * heightRatio;

        game->camera.zoom = newZoom;
    }

    // Limit camera to world
    {
        const f32 zoom = game->camera.zoom;
        const f32 width = GetScreenWidth() / zoom;
        const f32 height = GetScreenHeight() / zoom;
        const f32 mapWidth = game->map.width * game->map.tileWidth;
        const f32 mapHeight = game->map.height * game->map.tileHeight;
        game->camera.target.x = Clamp(game->camera.target.x,
                                      width * 0.5f - game->map.tileWidth,
                                      mapWidth - width + game->map.tileWidth + width * 0.5f);
        game->camera.target.y = Clamp(game->camera.target.y,
                                      height * 0.5f - game->map.tileHeight,
                                      mapHeight - height + game->map.tileHeight + height * 0.5f);
    }


    BzTileMap *map = &game->map;

    BzTile tileX = 0, tileY = 0;
    bzTileMapPosToTile(map, input->mouseWorld, &tileX, &tileY);

    if (IsKeyPressed(input->mapping.backBtn)) {
        ecs_remove_all(ECS, Selected);
        resetInputState(input);
        return;
    }

    const MouseButton primaryBtn   = input->mapping.primaryBtn;
    const MouseButton secondaryBtn = input->mapping.secondaryBtn;

    switch (input->state) {
        case INPUT_NONE: {
            inputPrimaryAction(game, input);
            break;
        }
        case INPUT_BUILDING: {
            BZ_ASSERT(input->building);
            BzTile sizeX = 0, sizeY = 0;
            getBuildingSize(input->building, &sizeX, &sizeY);
            bool canPlace = canPlaceBuilding(&game->map, input->building, tileX, tileY);
            if (canPlace && isInputBtnDown(input, primaryBtn)) {
                placeBuilding(&game->map, input->building, tileX, tileY);
            }
            input->buildingCanPlace = canPlace;
            input->buildingPos = (TilePosition) {tileX, tileY};
            input->buildingSize = (TileSize) {sizeX, sizeY};
            break;
        }
        case INPUT_SELECTED_UNITS: {
            inputPrimaryAction(game, input);
            if (input->state != INPUT_SELECTED_UNITS) break;
            inputUnitAction(game, input);

            /*

            i32 selectedCount = ecs_query_entity_count(input->queries.selected);
            if (selectedCount > 1 && isInputBtnJustDragged(input, secondaryBtn)) {
                // TODO: For click it should just move them
                i32 numUnits = selectedCount;
                f32 unitSpacing = 3.5f;
                bzArrayClear(input->unitPositions);
                placeUnits(numUnits, unitSpacing,
                    input->mouseDownWorld, input->mouseWorld,
                    map, &input->unitPositions);

                BZ_ASSERT(bzArraySize(input->unitPositions) == numUnits);
                i32 unitPosIdx = 0;

                ecs_defer_begin(ECS);
                ecs_defer_end(ECS);

                ecs_iter_t it = ecs_query_iter(ECS, input->queries.selected);
                ecs_defer_begin(ECS);
                while (ecs_iter_next(&it)) {
                    Position *pos = ecs_field(&it, Position, 1);
                    for (i32 i = 0; i < it.count; i++) {
                        ecs_entity_t entity = it.entities[i];

                        Position target = bzArrayGet(input->unitPositions, unitPosIdx);
                        unitPosIdx++;

                        Path path = {NULL, 0};
                        pathfindAStar(&(PathfindingDesc) {
                                .start=pos[i],
                                .target=target,
                                .map=map,
                                .outPath=&path,
                                .pool=game->pools.pathData,
                                .alloc=&game->stackAlloc
                        });
                        if (!path.paths) continue;
                        ecs_set_ptr(ECS, entity, Path, &path);

                    }
                }
                ecs_defer_end(ECS);
            } else if (isInputBtnJustDown(input, secondaryBtn)) {
                ecs_defer_begin(ECS);
                iterateSelectedUnits(input->queries.selected, iterRemovePaths);
                ecs_defer_end(ECS);

                ecs_iter_t it = ecs_query_iter(ECS, input->queries.selected);

                ecs_defer_begin(ECS);
                while (ecs_iter_next(&it)) {
                    Position *pos = ecs_field(&it, Position, 1);
                    for (i32 i = 0; i < it.count; i++) {
                        ecs_entity_t entity = it.entities[i];
                        ecs_remove(ECS, entity, Path);

                        Path path = {NULL, 0};
                        pathfindAStar(&(PathfindingDesc) {
                                .start=pos[i],
                                .target=input->mouseWorld,
                                .map=map,
                                .outPath=&path,
                                .pool=game->pools.pathData,
                                .alloc=&game->stackAlloc
                        });
                        if (!path.paths) continue;
                        ecs_set_ptr(ECS, entity, Path, &path);

                    }
                }
                ecs_defer_end(ECS);
            }
            */
            break;
        }
        case INPUT_SELECTED_OBJECT: {
            inputPrimaryAction(game, input);
            if (input->state != INPUT_SELECTED_OBJECT) break;
            break;
        }
        case INPUT_SELECTED_BUILDING: {
            inputPrimaryAction(game, input);
            if (input->state != INPUT_SELECTED_BUILDING) break;
            break;
        }

    }
}

void drawPlayerInputUIGround() {
    const Game       *game  = ecs_singleton_get_mut(ECS, Game);
    const InputState *input = ecs_singleton_get_mut(ECS, InputState);

    const MouseButton primaryBtn   = input->mapping.primaryBtn;

    const i32 selectedUnitCount = ecs_count_id(ECS, Selected);
    switch (input->state) {
        case INPUT_BUILDING: {
            const Color placeColor = input->buildingCanPlace ?
                               (Color) {0, 255, 0, 200} :
                               (Color) {255, 0, 0, 200};

            const BzTile width = game->map.tileWidth;
            const BzTile height = game->map.tileHeight;
            DrawRectangleLines(input->buildingPos.x * width,
                               input->buildingPos.y * height,
                               input->buildingSize.sizeX * width,
                               input->buildingSize.sizeY * height, placeColor);
            break;
        }
        case INPUT_SELECTED_UNITS: {
            //BZ_ASSERT(selectedUnitCount);
            /*
            if (selectedUnitCount > 1 && isInputBtnDragged(input, primaryBtn)) {
                i32 numUnits = selectedUnitCount;
                f32 unitSpacing = 3.5f;
                bzArrayClear(input->unitPositions);
                placeUnits(numUnits, unitSpacing, input->mouseDownWorld, input->mouseWorld,
                           map, &input->unitPositions);

                BZ_ASSERT(bzArraySize(input->unitPositions) == numUnits);
                bzArrayFor(input->unitPositions, i) {
                    Position pos = bzArrayGet(input->unitPositions, i);
                    DrawCircle(pos.x, pos.y, 2.0f, ORANGE);
                }
            }
            */
            break;
        }
        default: break;
    }
    ecs_iter_t it = ecs_query_iter(ECS, input->queries.selected);
    rlSetLineWidth(2.0f);
    while (ecs_query_next(&it)) {
        Position *pos  = ecs_field(&it, Position, 1);
        Size     *size = ecs_field(&it, Size, 2);
        for (i32 i = 0; i < it.count; i++) {
            f32 radius = size[i].x;
            if (size[i].y > radius)
                radius = size[i].y;
            radius *= 0.5f;
            DrawCircleLines(pos[i].x, pos[i].y, radius, GREEN);
        }
    }

}

void drawPlayerInputUI() {
    const Game       *game  = ecs_singleton_get_mut(ECS, Game);
    const InputState *input = ecs_singleton_get_mut(ECS, InputState);

    const MouseButton primaryBtn   = input->mapping.primaryBtn;

    if (isInputBtnDragged(input, primaryBtn)) {
        const Rectangle area = input->pickArea;
        DrawRectangleLines(area.x, area.y, area.width, area.height, RED);
    }

    switch (input->cursor) {
        case CURSOR_COLLECT_WOOD: {
            const Vector2 point = input->mouseWorld;
            DrawCircle(point.x, point.y, 2.0f, RED);
            DrawText("Collect wood", point.x, point.y, 10.0f, RED);
            break;
        }
        default: break;
    }

}


Rectangle calculateEntityBounds(Position pos, Size size) {
    return (Rectangle) {
        pos.x - size.x * 0.5f,
        pos.y - size.x * 0.5f,
        size.x, size.y
    };
}

bool getEntityBounds(ecs_entity_t entity, Position *outPos, Size *outSize, Rectangle *outBounds) {
    if (!ecs_is_alive(ECS, entity))
        return false;
    const Position *pos = ecs_get(ECS, entity, Position);
    if (!pos)
        return false;
    const Size *size = ecs_get(ECS, entity, Size);
    if (!size)
        return false;

    if (outPos) {
        *outPos = *pos;
    }
    if (outSize) {
        *outSize = *size;
    }
    if (outBounds) {
        *outBounds = (Rectangle) {
                pos->x - size->x * 0.5f,
                pos->y - size->y * 0.5f,
                size->x, size->y
        };
    }
    return true;
}

ecs_entity_t queryEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag) {
    BzSpatialGridIter it = bzSpatialGridIter(entityGrid, point.x, point.y, 0.0f, 0.0f);
    f32 closestDst = INFINITY;
    ecs_entity_t closest = 0;
    while (bzSpatialGridQueryNext(&it)) {
        ecs_entity_t entity = *(ecs_entity_t *) it.data;
        if (!ecs_is_alive(ECS, entity)) continue;
        if (!ecs_has_id(ECS, entity, Selectable)) continue;
        if (!ecs_has_id(ECS, entity, tag)) continue;
        Vector2 pos;
        Rectangle bounds;
        if (!getEntityBounds(entity, &pos, NULL, &bounds)) continue;

        if (!CheckCollisionPointRec(point, bounds)) continue;

        f32 curDst = Vector2Distance(point, pos);
        if (closestDst > curDst) {
            closestDst = curDst;
            closest = entity;
        }
    }
    return closest;
}

bool pickEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag) {
    ecs_remove_all(ECS, Selected);
    const ecs_entity_t entity = queryEntity(entityGrid, point, tag);
    if (entity) {
        ecs_add(ECS, entity, Selected);
        return true;
    }
    return false;
}

void pickUnits(BzSpatialGrid *entityGrid, Rectangle area) {
    ecs_remove_all(ECS, Selected);
    BzSpatialGridIter it = bzSpatialGridIter(entityGrid, area.x, area.y, area.width, area.height);
    while (bzSpatialGridQueryNext(&it)) {
        ecs_entity_t entity = *(ecs_entity_t *) it.data;
        if (!ecs_has_id(ECS, entity, ecs_id(Unit))) continue;
        Rectangle bounds;
        if (!getEntityBounds(entity, NULL, NULL, &bounds)) continue;

        if (!CheckCollisionRecs(area, bounds)) continue;
        ecs_add(ECS, entity, Selected);
    }
}



static bool isUnitObstructed(f32 x, f32 y, BzTileMap *map) {
    return bzTileMapHasCollision(map, x / map->tileWidth, y / map->tileHeight);
}

static bool canPlaceUnit(Vector2 pos, f32 space, BzTileMap *map) {
    for (i32 y = -1; y <= 1; y++) {
        for (i32 x = -1; x <= 1; x++) {
            if (isUnitObstructed(pos.x + x * space, pos.y + y * space, map))
                return false;
        }
    }

    return true;
}

void placeUnits(i32 numUnits, f32 unitSpacing, Vector2 start, Vector2 end, BzTileMap *map, Vector2 **outPlaces) {
    f32 angle = Vector2Angle(start, end);

    f32 lineLength = Vector2Distance(start, end);

    Vector2 pos = Vector2Zero();
    pos.x = unitSpacing;

    for (i32 i = 0; i < numUnits; i++) {
        if (pos.x + unitSpacing * 2.0f > lineLength) {
            pos.x = unitSpacing;
            pos.y += unitSpacing * 2.0f;
        }
        Vector2 unitPos = Vector2Add(start, Vector2Rotate(pos, angle));
        if (!canPlaceUnit(unitPos, 4.0f, map)) {
            i--;
        } else {
            bzArrayPush(*outPlaces, unitPos);
        }
        pos.x += unitSpacing * 2.0f;
    }
}

static void iterateSelectedUnits(ecs_query_t *query, void (*fn)(ecs_entity_t entity, Position *pos, Size *size)) {
    ecs_iter_t it = ecs_query_iter(ECS, query);
    while (ecs_iter_next(&it)) {
        Position *pos  = ecs_field(&it, Position, 1);
        Size     *size = ecs_field(&it, Size, 2);

        for (i32 i = 0; i < it.count; i++) {
            ecs_entity_t entity = it.entities[i];

            fn(entity, pos + i, size + i);
        }
    }
}
static void iterRemovePaths(ecs_entity_t entity, Position *pos, Size *size) {
    ecs_remove(ECS, entity, Path);
}