#include "systems.h"
#include "game_state.h"
#include "buildings.h"
#include "pathfinding.h"
#include <rlImGui.h>
#include <raymath.h>

bool getEntityBounds(ecs_entity_t entity, Position *outPos, Size *outSize, Rectangle *outBounds);
void pickEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t **outEntities);
void pickEntities(BzSpatialGrid *entityGrid, Rectangle area, ecs_entity_t **outEntities);

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

void drawSelectionCircles(ecs_entity_t **selected);

static bool isInputDragged(InputState *input) {
    return IsMouseButtonDown(input->LMB && input->mouseDownElapsed > input->CLICK_LIMIT);
}
static bool wasInputDragged(InputState *input) {
    return IsMouseButtonReleased(input->LMB) && input->mouseDownElapsed > input->CLICK_LIMIT;
}
static bool wasInputClicked(InputState *input) {
    return IsMouseButtonReleased(input->LMB);
}

void updatePlayerInput(ecs_iter_t *it) {
    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);
    if (IsKeyDown(KEY_W)) game->camera.target.y -= 20;
    if (IsKeyDown(KEY_S)) game->camera.target.y += 20;
    if (IsKeyDown(KEY_A)) game->camera.target.x -= 20;
    if (IsKeyDown(KEY_D)) game->camera.target.x += 20;

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


    game->camera.zoom += ((float) GetMouseWheelMove() * 0.05f);
    BzTileMap *map = &game->map;

    if (IsMouseButtonPressed(input->LMB)) {
        input->mouseDown = GetMousePosition();
        input->mouseDownWorld = GetScreenToWorld2D(input->mouseDown, game->camera);
        input->mouseDownElapsed = 0;
    } else if (IsMouseButtonDown(input->LMB)) {
        input->mouseDownElapsed += dt;
    }

    Vector2 worldPos = GetScreenToWorld2D(GetMousePosition(), game->camera);
    BzTile tileX = 0, tileY = 0;
    bzTileMapPosToTile(map, worldPos, &tileX, &tileY);
    switch (input->state) {
        case INPUT_NONE:
            if (wasInputDragged(input)) {
                if (bzArraySize(input->entities) > 0) {
                    input->state = INPUT_SELECTED_UNITS;
                    break;
                }

            } else if (wasInputClicked(input)) {
                // Click

                // 1. Entity
                if (bzArraySize(input->entities) > 0) {
                    input->state = INPUT_SELECTED_UNITS;
                    break;
                }

                // 2. Object

                // 3. Building
            }
            break;
        case INPUT_BUILDING:
            if (IsKeyPressed(input->ESC) ||
                IsMouseButtonPressed(input->RMB) ||
                input->building == 0) {
                input->state = INPUT_NONE;
                input->building = 0;
                break;
            }
            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 && IsMouseButtonDown(input->LMB)) {
                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:
            if (IsKeyPressed(input->ESC) || IsMouseButtonPressed(input->RMB)) {
                bzArrayClear(input->entities);
                input->state = INPUT_NONE;
                break;
            }
            if (bzArraySize(input->entities) > 1 && wasInputDragged(input)) {
                // TODO: For click it should just move them
                i32 numUnits = bzArraySize(input->entities);
                f32 unitSpacing = 3.5f;
                bzArrayClear(input->unitPositions);
                placeUnits(numUnits, unitSpacing, input->mouseDownWorld, worldPos,
                           map, &input->unitPositions);

                BZ_ASSERT(bzArraySize(input->unitPositions) == numUnits);

                bzArrayFor(input->unitPositions, i) {
                    Position target = bzArrayGet(input->unitPositions, i);
                    ecs_entity_t entity = bzArrayGet(input->entities, i);

                    const Position *start = ecs_get(ECS, entity, Position);
                    Path path = {NULL, 0};
                    pathfindAStar(&(PathfindingDesc) {
                            .start=*start,
                            .target=target,
                            .map=map,
                            .outPath=&path,
                            .pool=game->pools.pathData
                    });
                    if (!path.paths) continue;
                    ecs_set_ptr(ECS, entity, Path, &path);
                }
            } else if (wasInputClicked(input)) {
                bzArrayFor(input->entities, i) {
                    ecs_entity_t entity = bzArrayGet(input->entities, i);
                    ecs_remove(ECS, entity, Path);

                    const Position *start = ecs_get(ECS, entity, Position);
                    Path path = {NULL, 0};
                    pathfindAStar(&(PathfindingDesc) {
                            .start=*start,
                            .target=worldPos,
                            .map=map,
                            .outPath=&path,
                            .pool=game->pools.pathData
                    });
                    if (!path.paths) continue;
                    ecs_set_ptr(ECS, entity, Path, &path);
                }

            }
            break;
        case INPUT_SELECTED_OBJECT:
            break;
        case INPUT_SELECTED_BUILDING:
            break;
    }
}

void drawPlayerInputUI(ecs_iter_t *it) {
    Game *game = ecs_singleton_get_mut(ECS, Game);
    InputState *input = ecs_singleton_get_mut(ECS, InputState);
    BzTileMap *map = &game->map;
    Vector2 worldPos = GetScreenToWorld2D(GetMousePosition(), game->camera);

    switch (input->state) {
        case INPUT_NONE:
            if (isInputDragged(input)) {
                Vector2 start = input->mouseDownWorld;
                Vector2 end = worldPos;
                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;
                }
                Rectangle area = {start.x, start.y, end.x - start.x, end.y - start.y};
                DrawRectangleLines(area.x, area.y, area.width, area.height, RED);

                pickEntities(game->entityGrid, area, &input->entities);
            }
            break;
        case INPUT_BUILDING: {
            Color placeColor = input->buildingCanPlace ?
                               (Color) {0, 255, 0, 200} :
                               (Color) {255, 0, 0, 200};

            BzTile width = game->map.tileWidth;
            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: {
            if (bzArraySize(input->entities) > 1 && isInputDragged(input)) {
                i32 numUnits = bzArraySize(input->entities);
                f32 unitSpacing = 3.5f;
                bzArrayClear(input->unitPositions);
                placeUnits(numUnits, unitSpacing, input->mouseDownWorld, worldPos,
                           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;
        }
        case INPUT_SELECTED_OBJECT:
            break;
        case INPUT_SELECTED_BUILDING:
            break;
    }
    drawSelectionCircles(&input->entities);



}

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

void pickEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t **outEntities) {
    BZ_ASSERT(*outEntities);
    bzArrayClear(*outEntities);
    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;
        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;
        }
    }
    if (closest) {
        bzArrayPush(*outEntities, closest);
    }
}

void pickEntities(BzSpatialGrid *entityGrid, Rectangle area, ecs_entity_t **outEntities) {
    BZ_ASSERT(*outEntities);
    bzArrayClear(*outEntities);

    BzSpatialGridIter it = bzSpatialGridIter(entityGrid, area.x, area.y, area.width, area.height);
    while (bzSpatialGridQueryNext(&it)) {
        ecs_entity_t entity = *(ecs_entity_t *) it.data;
        Rectangle bounds;
        if (!getEntityBounds(entity, NULL, NULL, &bounds)) continue;

        if (!CheckCollisionRecs(area, bounds)) continue;
        bzArrayPush(*outEntities, entity);
    }
}



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

void drawSelectionCircles(ecs_entity_t **selected) {
    bzArrayFor(*selected, i) {
        ecs_entity_t entity = bzArrayGet(*selected, i);
        Position pos = {0};
        if (!getEntityBounds(entity, &pos, NULL, NULL))
            continue;
        DrawCircleLines(pos.x, pos.y, 4.5f, GREEN);


    }
}