#include "ai_actions.h"
#include "game_state.h"
#include "components.h"
#include "systems/systems.h"
#include "buildings.h"

#include <raymath.h>

float shortestArc(float a, float b) {
    if (fabs(b - a) < M_PI)
        return b - a;
    if (b > a)
        return b - a - M_PI * 2.0f;
    return b - a + M_PI * 2.0f;
}

BzBTStatus aiMoveTo(AIBlackboard *data, f32 dt) {
    Game *game = ecs_singleton_get_mut(ECS, Game);
    const Vector2 pos = *ecs_get(ECS, data->entity, Position);
    const Vector2 target = data->moveToPos;
    f32 dst = Vector2Distance(pos, target);
    if (dst < data->proximity) {
        ecs_remove(ECS, data->entity, Path);
        return BZ_BT_SUCCESS;
    }
    if (!ecs_has(ECS, data->entity, Path)) {
        entitySetPath(data->entity, target, game);
    }
    if (ecs_has(ECS, data->entity, Orientation)) {
        Orientation *orientation = ecs_get_mut(ECS, data->entity, Orientation);
        f32 currentAngle = *orientation;
        f32 targetAngle = Vector2Angle(pos, target);
        f32 dif = shortestArc(currentAngle, targetAngle);
        dif = Clamp(dif, -1, 1) * dt * 10;
        *orientation += dif;
        *orientation = fmodf(*orientation + 180.0f, 360.0f) - 180.0f;
    }
    return BZ_BT_RUNNING;
}

BzBTStatus aiResetElapsed(AIBlackboard *data, f32 dt) {
    data->elapsed = 0.0f;
    return BZ_BT_SUCCESS;
}
BzBTStatus aiFindNextHarvestable(AIBlackboard *data, f32 dt) {
    ecs_entity_t harvestTarget = data->as.worker.harvestTarget;
    if (ecs_is_alive(ECS, harvestTarget)) {
        BZ_ASSERT(ecs_has_id(ECS, harvestTarget, Harvestable));
        // Target still alive, no need to find next harvestable
        data->moveToPos = data->as.worker.harvestPos;
        return BZ_BT_SUCCESS;
    }

    BZ_ASSERT(ecs_has(ECS, data->entity, Worker));
    Worker *worker = ecs_get_mut(ECS, data->entity, Worker);
    Game *game = ecs_singleton_get_mut(ECS, Game);

    ResourceType harvestType = data->as.worker.harvestType;

    // Perform spatial search

    ecs_entity_t closest = 0;
    f32 closestDst = INFINITY;
    Vector2 closestPos = Vector2Zero();

    const f32 range = 20.0f;
    f32 hRange = range * 0.5f;
    Vector2 pos = data->as.worker.harvestPos; // Last know harvest pos
    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_id(ECS, entity, Harvestable) ||
            !ecs_has(ECS, entity, Resource) ||
            !ecs_has(ECS, entity, Position))
            continue;
        Resource resource = *ecs_get(ECS, entity, Resource);
        Position resPos  = *ecs_get(ECS, entity, Position);
        if (resource.type != harvestType) continue;

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

    if (closest) {
        data->as.worker.harvestTarget = closest;
        data->as.worker.harvestPos = closestPos;
        data->moveToPos = closestPos;
        return BZ_BT_SUCCESS;
    }

    return BZ_BT_FAIL;
}
BzBTStatus aiFindNearestStorage(AIBlackboard *data, f32 dt) {
    ecs_filter_t *storageFilter = ecs_filter(ECS, {
        .terms = {{ecs_id(Position)}, {ecs_id(Storage)}},
    });
    ecs_iter_t it = ecs_filter_iter(ECS, storageFilter);

    const Vector2 pos = *ecs_get(ECS, data->entity, Position);

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

        for (i32 i = 0; i < it.count; i++) {
            f32 dst = Vector2Distance(pos, closestPos);
            if (closestDst == INFINITY || dst < closestDst) {
                closest    = it.entities[i];
                closestDst = dst;
                closestPos = storagePos[i];
            }
        }
    }
    ecs_filter_fini(storageFilter);

    if (!closest) {
        return BZ_BT_FAIL;
    }

    data->as.worker.depositTarget = closest;
    data->moveToPos = getPositionNearBuilding(closest, pos);

    return BZ_BT_SUCCESS;
}
BzBTStatus aiHarvestRes(AIBlackboard *data, f32 dt) {
    BZ_ASSERT(ecs_has(ECS, data->entity, Worker));
    Worker *worker = ecs_get_mut(ECS, data->entity, Worker);
    if (worker->carry >= worker->carryCapacity)
        return BZ_BT_FAIL;

    ecs_entity_t harvestTarget = data->as.worker.harvestTarget;
    if (!ecs_is_alive(ECS, harvestTarget))
        return BZ_BT_FAIL;

    BZ_ASSERT(ecs_has_id(ECS, harvestTarget, Harvestable));
    if (data->elapsed < worker->collectSpeed) {
        data->elapsed += dt;
        return BZ_BT_RUNNING;
    }
    data->elapsed = 0;

    // Collect
    i32 spareCapacity = worker->carryCapacity - worker->carry;
    BZ_ASSERT(spareCapacity >= 0);
    i32 collected = harvestEvent(harvestTarget, (HarvestEvent) {
        .amount = BZ_MIN(1, spareCapacity)
    });
    worker->carry += collected;

    return BZ_BT_SUCCESS;
}
BzBTStatus aiDepositRes(AIBlackboard *data, f32 dt) {
    BZ_ASSERT(ecs_has(ECS, data->entity, Worker));
    Worker *worker = ecs_get_mut(ECS, data->entity, Worker);

    if (worker->carry == 0)
        return BZ_BT_SUCCESS;

    ecs_entity_t depositTarget = data->as.worker.depositTarget;

    if (!ecs_is_alive(ECS, depositTarget))
        return BZ_BT_FAIL;

    if (data->elapsed < worker->depositSpeed) {
        data->elapsed += dt;
        return BZ_BT_RUNNING;
    }

    depositEvent(depositTarget, (DepositEvent) {
        .amount = worker->carry
    });
    worker->carry = 0;

    return BZ_BT_SUCCESS;
}
BzBTStatus aiCarryCapacityFull(AIBlackboard *data) {
    BZ_ASSERT(ecs_has(ECS, data->entity, Worker));
    Worker *worker = ecs_get_mut(ECS, data->entity, Worker);
    if (worker->carry >= worker->carryCapacity)
        return BZ_BT_SUCCESS;
    return BZ_BT_FAIL;
}
BzBTStatus aiCarryCapacityEmpty(AIBlackboard *data) {
    BZ_ASSERT(ecs_has(ECS, data->entity, Worker));
    Worker *worker = ecs_get_mut(ECS, data->entity, Worker);
    if (worker->carry == 0)
        return BZ_BT_SUCCESS;
    return BZ_BT_FAIL;
}
