#include "systems.h"

#include "../game_state.h"
#include "../input.h"
#include "../building_factory.h"
#include "../pathfinding.h"

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

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

bool selectEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag, Player player);
void selectUnits(BzSpatialGrid *entityGrid, Rectangle area, Player player);

static bool selectedAnyHasID(ecs_query_t *query, ecs_entity_t id);

void addEntityToInspected(ecs_entity_t entity, Game *game);

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

i32 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 = BUILDING_NONE;
}

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};
        selectUnits(game->entityGrid, input->pickArea, game->player);
    }
    i32 selectedCount = ecs_query_entity_count(input->queries.selected);
    if (isInputBtnJustDragged(input, primaryBtn) && selectedCount > 0) {
        resetInputState(input);
        input->state = INPUT_SELECTED_UNITS;
    }
    if (IsKeyDown(KEY_LEFT_ALT) && isInputBtnJustUp(input, primaryBtn)) {
        ecs_entity_t entity = queryEntity(game->entityGrid, input->mouseDownWorld, ecs_id(GameEntity));
        if (entity) addEntityToInspected(entity, game);
    } else if (isInputBtnJustUp(input, primaryBtn)) {
        InputType type = input->state;
        if (selectEntity(game->entityGrid, input->mouseDownWorld, ecs_id(Unit), game->player))
            type = INPUT_SELECTED_UNITS;
        else if (selectEntity(game->entityGrid, input->mouseDownWorld, ecs_id(Harvestable), PLAYER_COUNT))
            type = INPUT_SELECTED_OBJECT;
        else if (selectEntity(game->entityGrid, input->mouseDownWorld, ecs_id(Building), game->player))
            type = INPUT_SELECTED_BUILDING;
        selectedCount = ecs_query_entity_count(input->queries.selected);
        if (selectedCount > 0) {
            resetInputState(input);
            input->state = type;
        }
    }

    if (selectedCount == 0)
        resetInputState(input);
}

typedef struct EntityPosPair {
    ecs_entity_t entity;
    Vector2 pos;
} EntityPosPair;
static int vec2CmpZero(const void *lhsData, const void *rhsData) {
    const Vector2 *lhs = lhsData;
    const Vector2 *rhs = rhsData;

    f32 dstL = Vector2DistanceSqr(Vector2Zero(), *lhs);
    f32 dstR = Vector2DistanceSqr(Vector2Zero(), *rhs);

    if (dstL < dstR) return -1;
    if (dstL > dstR) return 1;
    return 0;
}
static int entityPosPairCmpZero(const void *lhsData, const void *rhsData) {
    const EntityPosPair *lhs = lhsData;
    const EntityPosPair *rhs = rhsData;

    return vec2CmpZero(&lhs->pos, &rhs->pos);
}
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;
    bool isWorker = selectedAnyHasID(query, ecs_id(Worker));
    if (isWorker && (taskEntity = queryEntity(game->entityGrid, input->mouseWorld, ecs_id(Harvestable)))) {
        Resource resource = *ecs_get(ECS, taskEntity, Resource);
        switch (resource.type) {
            case RES_WOOD:
                input->cursor = CURSOR_COLLECT_WOOD;
                break;
            case RES_GOLD:
                input->cursor = CURSOR_COLLECT_GOLD;
                break;
            case RES_FOOD:
                input->cursor = CURSOR_FARM;
                break;
            default:;
        }
        if (isInputBtnJustUp(input, actionBtn)) {
            const f32 hRadius = 10.0f;
            const Vector2 mPos = input->mouseWorld;
            BzSpatialGridIter gridIt = bzSpatialGridIter(game->entityGrid,
                                                         mPos.x - hRadius, mPos.y - hRadius,
                                                         mPos.x + hRadius, mPos.y + hRadius);
            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++) {
                    ecs_entity_t entity = it.entities[i];

                    bool hasNext = false;
                    ecs_entity_t harvestEntity = 0;
                    do {
                        hasNext = bzSpatialGridQueryNext(&gridIt);
                        if (!hasNext) break;
                        harvestEntity = *(ecs_entity_t *) gridIt.data;
                        if (!ecs_has(ECS, harvestEntity, Resource))
                            continue;
                        const Resource *res = ecs_get(ECS, harvestEntity, Resource);
                        if (res->type != resource.type)
                            continue;
                        if (!ecs_has(ECS, harvestEntity, Harvestable))
                            continue;
                        Harvestable *harvestable = ecs_get_mut(ECS, harvestEntity, Harvestable);
                        if (harvestable->harvestCount >= harvestable->harvestLimit)
                            continue;
                        harvestable->harvestCount++;
                        break;
                    } while (hasNext);

                    if (!hasNext) break;
                    Position target = *ecs_get(ECS, harvestEntity, Position);
                    HitBox targetHB = *ecs_get(ECS, harvestEntity, HitBox);
                    target = entityGetCenter(target, targetHB);

                    f32 proximity = 4.0f;

                    Worker *worker = ecs_get_mut(ECS, entity, Worker);
                    worker->carryRes = resource.type;

                    setAIBehaviour(entity, game->BTs.workerHarvest, &(AIBlackboard) {
                            .as.worker = {
                                    .harvestType = resource.type,
                                    .harvestTarget = harvestEntity,
                                    .harvestPos = target,
                            },
                            .proximity = proximity,
                    });
                }
            }
            ecs_defer_end(ECS);
        }
        return;
    }

    // Unit place position
    Vector2 *positions = bzStackAlloc(&game->stackAlloc, sizeof(*positions) * numUnits);
    Vector2 start = Vector2Zero();
    Vector2 end = Vector2Zero();
    if (isInputBtnDragged(input, actionBtn) || isInputBtnJustDragged(input, actionBtn)) {
        start = input->mouseDownWorld;
        end = input->mouseWorld;
    } else {
        start = end = input->mouseWorld;
        f32 displace = 8 * numUnits;
        displace = BZ_MIN(displace, 50.0f);
        start.x -= displace;
        end.x += displace + 2;
    }
    i32 placedUnits = placeUnits(numUnits, 6.0f, start, end, map, positions);
    input->numUnits = placedUnits;
    input->unitPlacePos = positions;

    if (IsMouseButtonReleased(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);


        EntityPosPair *entities = bzStackAlloc(&game->stackAlloc, sizeof(*entities) * input->numUnits);
        ecs_iter_t it = ecs_query_iter(ECS, query);
        i32 unitIdx = 0;
        while (ecs_iter_next(&it)) {
            for (i32 i = 0; i < it.count; i++) {
                if (unitIdx >= input->numUnits)
                    break;
                const ecs_entity_t entity = it.entities[i];
                if (!ecs_has(ECS, entity, Position)) continue;
                entities[unitIdx++] = (EntityPosPair) {
                    .entity = entity,
                    .pos = *ecs_get(ECS, entity, Position)
                };
            }
        }

        qsort(entities, unitIdx, sizeof(*entities), entityPosPairCmpZero);
        qsort(input->unitPlacePos, unitIdx, sizeof(*input->unitPlacePos), vec2CmpZero);

        for (i32 i = 0; i < unitIdx; i++) {
            ecs_entity_t entity = entities[i].entity;
            setAIBehaviour(entity, game->BTs.moveTo, &(AIBlackboard) {
                    .moveToPos = positions[i],
                    .proximity = 1.0f,
            });
        }
    }

}

void updatePlayerInput() {
    f32 dt = GetFrameTime();

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

    input->numUnits = 0;
    input->unitPlacePos = NULL;

    const f32 maxZoom = 4.5f;
    const f32 minZoom = 0.9f;
    if (input->canUseKeyboard) {
        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++;

        if (IsKeyReleased(input->mapping.pauseBtn)) {
            setScreen(game, SCREEN_PAUSE_MENU);
        } else if (IsKeyReleased(input->mapping.backBtn)) {
            if (input->state == INPUT_NONE) {
                setScreen(game, SCREEN_PAUSE_MENU);
            } else {
                ecs_remove_all(ECS, Selected);
                resetInputState(input);
            }

        }

    }

    // https://gamedev.stackexchange.com/questions/9330/zoom-to-cursor-calculation
    float wheel = GetMouseWheelMove();
    if (wheel != 0.0f && input->canUseMouse) {
        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);
    }

    if (!input->canUseMouse)
        return;


    BzTileMap *map = &game->map;

    Vec2i tileXY = bzTileMapPosToTile(map, input->mouseWorld);
    BzTile tileX = tileXY.x;
    BzTile tileY = tileXY.y;

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

    i32 count = ecs_query_entity_count(input->queries.selected);
    //if (count > 0)
    //    input->state = INPUT_SELECTED_UNITS;

    switch (input->state) {
        case INPUT_NONE: {
            inputPrimaryAction(game, input);
            break;
        }
        case INPUT_BUILDING: {
            if (input->building <= BUILDING_NONE || input->building >= BUILDING_COUNT ||
                    isInputBtnJustUp(input, input->mapping.secondaryBtn)) {
                input->state = INPUT_NONE;
                resetInputState(input);
                return;
            }
            BzTile sizeX = 0, sizeY = 0;
            getBuildingSize(input->building, &sizeX, &sizeY);
            bool canPlace = canPlaceBuilding(game, input->building, tileX, tileY);
            if (canPlace && isInputBtnDown(input, primaryBtn)) {
                placeBuilding(game, input->building, tileX, tileY, game->player);
                // Update player resources
                i32 cost[RES_COUNT] = {0};
                getBuildingCost(input->building, cost);
                game->playerResources[game->player].wood -= cost[RES_WOOD];
                game->playerResources[game->player].gold -= cost[RES_GOLD];
                game->playerResources[game->player].food -= cost[RES_FOOD];
            }
            input->buildingCanPlace = canPlace;
            input->buildingPos = (Vec2i) {tileX, tileY};
            input->buildingSize = (Vec2i) {sizeX, sizeY};
            break;
        }
        case INPUT_SELECTED_UNITS: {
            inputPrimaryAction(game, input);
            if (input->state != INPUT_SELECTED_UNITS) break;
            inputUnitAction(game, input);
            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;

    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.x * width,
                               input->buildingSize.y * height, placeColor);
            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);
        HitBox   *hitbox = ecs_field(&it, HitBox , 2);
        for (i32 i = 0; i < it.count; i++) {
            ecs_entity_t entity = it.entities[i];
            f32 radius = BZ_MAX(hitbox[i].width, hitbox[i].height);
            Position center = entityGetCenter(pos[i], hitbox[i]);
            radius *= 0.8f;
            const f32 lineThickness = 1.0f;
            if (ecs_has(ECS, entity, Building)) {
                const f32 padding = 2.0f;
                Rectangle bounds = {
                        pos[i].x + hitbox[i].x - padding,
                        pos[i].y - hitbox[i].height - padding,
                        hitbox[i].width + padding * 2,
                        hitbox[i].height + padding * 2,
                };
                DrawRectangleLinesEx(bounds, lineThickness, GREEN);
            } else {
                DrawRing(center, radius, radius + lineThickness, 0, 360, 12, GREEN);
            }
        }
    }

    if (input->unitPlacePos) {
        for (i32 i = 0; i < input->numUnits; i++) {
            DrawCircleV(input->unitPlacePos[i], 2.0f, RED);
        }
    }

}

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

    Vector2 point = input->mouseWorld;
    Rectangle texRect = {0, 0, 0, 0};
    switch (input->cursor) {
        case CURSOR_COLLECT_WOOD:
            texRect = getTextureRect(getItemTile(ITEM_AXE));
            break;
        case CURSOR_COLLECT_GOLD:
            texRect = getTextureRect(getItemTile(ITEM_PICKAXE));
            break;
        case CURSOR_FARM:
            texRect = getTextureRect(getItemTile(ITEM_SYTHE));
            break;
        case CURSOR_ATTACK:
            texRect = getTextureRect(getItemTile(ITEM_CUTLASS));
            break;
        default: break;
    }

    if (texRect.width != 0 && texRect.height != 0) {
        Texture tiles = game->tileset.tiles;
        f32 size = 35.0 / game->camera.zoom;
        Rectangle dst = {
                point.x - size * 0.5f,
                point.y - size * 0.5f,
                size, size
        };
        point.y -= texRect.height;

        DrawTexturePro(tiles, texRect, dst, Vector2Zero(), 0.0f, WHITE);
    }

}

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 hitbox;
        if (!entityGetHitBox(entity, &pos, &hitbox)) continue;

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

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

bool selectEntity(BzSpatialGrid *entityGrid, Vector2 point, ecs_entity_t tag, Player player) {
    ecs_remove_all(ECS, Selected);
    const ecs_entity_t entity = queryEntity(entityGrid, point, tag);
    if (!entity) return false;
    if (player != PLAYER_COUNT) {
        if (!ecs_has(ECS, entity, Owner)) return false;
        Owner owner = *ecs_get(ECS, entity, Owner);
        if (player != owner.player) return false;
    }
    ecs_add(ECS, entity, Selected);
    return true;
}

void selectUnits(BzSpatialGrid *entityGrid, Rectangle area, Player player) {
    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 (player != PLAYER_COUNT) {
            if (!ecs_has(ECS, entity, Owner)) continue;
            Owner owner = *ecs_get(ECS, entity, Owner);
            if (owner.player != player) continue;
        }
        if (!ecs_has_id(ECS, entity, ecs_id(Unit))) continue;
        Rectangle hitbox;
        if (!entityGetHitBox(entity, NULL, &hitbox)) continue;

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

static bool selectedAnyHasID(ecs_query_t *query, ecs_entity_t id) {
    ecs_iter_t it = ecs_query_iter(ECS, query);
    while (ecs_iter_next(&it)) {
        for (i32 i = 0; i < it.count; i++) {
            ecs_entity_t entity = it.entities[i];
            if (ecs_has_id(ECS, entity, id)) {
                ecs_iter_fini(&it);
                return true;
            }
        }
    }
    return false;
}

void addEntityToInspected(ecs_entity_t entity, Game *game) {
    bool alreadyInspecting = false;
    for (i32 i = 0; i < bzArraySize(game->debug.inspecting); i++) {
        if (game->debug.inspecting[i] == entity) {
            alreadyInspecting = true;
            break;
        }
    }
    if (!alreadyInspecting)
        bzArrayPush(game->debug.inspecting, entity);
}

static bool isUnitObstructed(f32 x, f32 y, BzTileMap *map) {
    return bzTileMapHasAnyCollision(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;
}
static bool unitWithinMap(Vector2 pos, BzTileMap *map) {
    f32 x = pos.x / map->tileWidth;
    f32 y = pos.y / map->tileHeight;

    return !(x < 0 || y < 0 || x >= map->width || y >= map->height);
}
i32 placeUnits(i32 numUnits, f32 unitSpacing, Vector2 start, Vector2 end, BzTileMap *map, Vector2 *outPlaces) {
    i32 outIdx = 0;
    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));
        bool withinMap = unitWithinMap(unitPos, map);
        if (canPlaceUnit(unitPos, 4.0f, map) && withinMap) {
            outPlaces[outIdx++] = unitPos;
        } else if (withinMap) {
            i--;
        }
        pos.x += unitSpacing * 2.0f;
    }
    return outIdx;
}

static void iterateSelectedUnits(ecs_query_t *query, void (*fn)(ecs_entity_t entity, Position *pos)) {
    ecs_iter_t it = ecs_query_iter(ECS, query);
    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];

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