#include "unit_ai.h"
#include "unit_actions.h"
#include "game_state.h"
#include "buildings.h"

#include <raymath.h>

static ecs_entity_t findNearestStorage(Position pos, ResourceType type, Position *outPos) {
    ecs_filter_t *storageFilter = ecs_filter(ECS, {
        .terms = {{ecs_id(Storage)}, {ecs_id(Position)} }
    });
    ecs_iter_t it = ecs_filter_iter(ECS, storageFilter);

    ecs_entity_t closest = 0;
    f32          closestDst = INFINITY;
    Position     closestPos = Vector2Zero();
    while (ecs_filter_next(&it)) {
        Storage  *storage    = ecs_field(&it, Storage, 1);
        Position *storagePos = ecs_field(&it, Position, 2);

        for (i32 i = 0; i < it.count; i++) {
            f32 dst = Vector2Distance(pos, storagePos[i]);
            if (storage[i].capacity[type] && dst < closestDst) {
                closest    = it.entities[i];
                closestDst = dst;
                closestPos = storagePos[i];
            }
        }
    }
    ecs_filter_fini(storageFilter);
    if (outPos) *outPos = closestPos;
    return closest;
}
static ecs_entity_t findNearestResource(Game *game, Position pos, ResourceType type, f32 range) {
    ecs_entity_t closest    = 0;
    f32          closestDst = 10000.0f;

    f32 hRange = range * 0.5f;
    Rectangle area = {pos.x - hRange, pos.y + hRange, hRange, hRange};

    BzSpatialGridIter it = bzSpatialGridIter(game->entityGrid, area.x, area.y, area.width, area.height);

    while (bzSpatialGridQueryNext(&it)) {
        ecs_entity_t entity = *(ecs_entity_t *) it.data;
        if (!ecs_is_alive(ECS, entity)) continue;
        if (!ecs_has(ECS, entity, Resource) || !ecs_has(ECS, entity, Position))
            continue;
        const Resource *resource = ecs_get(ECS, entity, Resource);
        const Position *resPos   = ecs_get(ECS, entity, Position);
        if (resource->type != type) continue;

        f32 dst = Vector2Distance(pos, *resPos);
        if (dst < closestDst) {
            closest    = entity;
            closestDst = dst;
        }
    }

    return closest;
}

void aiWorkerBuild(ecs_entity_t entity, UnitAI *unitAI, Game *game) {

}
void aiWorkerHarvestUpdate(ecs_entity_t entity, UnitAI *unitAI, Game *game) {
    const Action *action = unitAI->action;
    if (action) return;

    AIWorkerHarvest worker = unitAI->as.workerHarvest;

    Position workerPos = *ecs_get(ECS, entity, Position);

    // Finished all tasks, repeat
    ecs_entity_t target    = worker.target;
    Position     targetPos = worker.targetPosition;
    if (!ecs_is_alive(ECS, target)) {
        // Find new resource
        target = findNearestResource(game, targetPos, worker.resource, 20.0f);
        if (!target) return;

        worker.target = target;
        worker.targetPosition = *ecs_get(ECS, target, Position);
    }
    // Find the closest warehouse
    Position warehousePos = Vector2Zero();
    ecs_entity_t warehouse = findNearestStorage(workerPos, worker.resource,
                                                &warehousePos);
    if (!warehouse) {
        return;
    }
    warehousePos = getPositionNearBuilding(warehouse, targetPos);

    const f32 proximityThreshold = 8.0f;

    addAction(entity, game, &(const Action) {
        .type = ACTION_MOVE_TO,
        .as.moveTo.target = targetPos,
        .as.moveTo.proximityThreshold = proximityThreshold,
    });
    addAction(entity, game, &(const Action) {
        .type = ACTION_COLLECT_RESOURCE,
        .as.collectResource.entity = target
    });
    addAction(entity, game, &(const Action) {
        .type = ACTION_MOVE_TO,
        .as.moveTo.target = warehousePos,
        .as.moveTo.proximityThreshold = proximityThreshold,
    });
    addAction(entity, game, &(const Action) {
       .type = ACTION_DEPOSIT_RESOURCE,
       .as.depositResource.entity = warehouse
    });

}

void updateUnitAI(ecs_entity_t entity, UnitAI *unitAI, Game *game) {
    //if (unitAI->onUpdate) unitAI->onUpdate(entity, unitAI, game);
    aiWorkerHarvestUpdate(entity, unitAI, game);
}
void setUnitAI(ecs_entity_t entity, Game *game, const UnitAI *unitAI) {
    BZ_ASSERT(unitAI);
    if (ecs_has(ECS, entity, UnitAI)) {
        UnitAI *entityAI = ecs_get_mut(ECS, entity, UnitAI);
        if (entityAI->onRemove) entityAI->onRemove(entity, entityAI, game);
    }
    clearActions(entity, game);

    UnitAI ai = *unitAI;
    if (unitAI->onSet) unitAI->onSet(entity, &ai, game);
    ecs_set_ptr(ECS, entity, UnitAI, &ai);
    if (unitAI->action != NULL)
        addAction(entity, game, unitAI->action);
}