#include "map_init.h"

#include <flecs.h>
#include <stdio.h>

#include "building_factory.h"
#include "components.h"
#include "entity_factory.h"
#include "game_state.h"
#include "map_layers.h"
#include "utils.h"
#include "systems/systems.h"
#include "raymath.h"

i32 getSwarmWaypointIdx(u32 id) {
    char buf[4];
    for (i32 i = 0; i < MAX_SWARM_WAYPOINTS; i++) {
        snprintf(buf, sizeof(buf), "p%d", i);
        if (id == bzStringDefaultHash(buf))
            return i;
    }
    return -1;
}

bool initGameObjectsLayer(BzTileMap *map, BzTileObjectGroup *objectGroup) {
    Game *game = ecs_singleton_get_mut(ECS, Game);
    game->swamNumWaypoints = 0;
    for (i32 i = 0; i < objectGroup->objectCount; i++) {
        BzTileObject object = objectGroup->objects[i];
        if (bzStringDefaultHash("camera") == object.id) {
            game->camera.target.x = object.shape.x;
            game->camera.target.y = object.shape.y;
        }
        i32 swarmIdx = getSwarmWaypointIdx(object.id);
        if (swarmIdx != -1) {
            game->swarmWaypoints[swarmIdx] = (Vector2) {
                object.shape.x,
                object.shape.y,
            };
            game->swamNumWaypoints = BZ_MAX(game->swamNumWaypoints, swarmIdx + 1);
        }
        if (bzStringDefaultHash("spawn") == object.id) {
            game->swarmSpawn.x = object.shape.x;
            game->swarmSpawn.y = object.shape.y;
        }
    }
    return true;
}



bool initEntityObjectsLayer(BzTileMap *map, BzTileObjectGroup *objectGroup) {
    Game *game = ecs_singleton_get_mut(ECS, Game);
    BzTileset *objectTileset = bzTileObjectGroupGetTileset(&game->map, objectGroup);
    if (!objectTileset) return true;

    for (i32 i = 0; i < objectGroup->objectCount; i++) {
        //if (i == 1) break;
        BzTileObject object = objectGroup->objects[i];
        Position pos = (Position) { object.shape.x, object.shape.y };
        BzTileID gid = bzTilesetGetTileID(objectTileset, object.gid);
        BzTileID baseGid = getTileBase(gid);
        Player player = PLAYER_RED;
        if (baseGid != gid)
            player = PLAYER_BLUE;
        EntityType entity = getTileEntity(baseGid);
        switch (entity) {
            case ENTITY_GOBLIN:
                break;
            case ENTITY_MAGE:
                break;
            case ENTITY_ORC:
                break;
            case ENTITY_SOLDIER:
                entityCreateSoldier(pos, player, game);
                break;
            case ENTITY_WARRIOR:
                break;
            case ENTITY_WORKER:
                entityCreateWorker(pos, player, game);
                break;
            default:
                BZ_ASSERT(false);
                break;
        }
    }
    return true;

}

bool initBuildingsLayer(BzTileMap *map, BzTileLayer *layer) {
    Game *game = ecs_singleton_get_mut(ECS, Game);

    BzTileLayer *ownershipLayer = layer;
    BzTileLayer *buildingLayer = bzTileMapGetLayer(map, LAYER_BUILDINGS);
    BzTileset *buildingTileset = bzTileLayerGetTileset(map, buildingLayer);
    BzTile *data = layer->data;
    BzTile *buildingData = buildingLayer->data;
    BZ_ASSERT(ownershipLayer->tilesetIdx == buildingLayer->tilesetIdx);

    for (i32 y = 0; y < layer->height; y++) {
        for (i32 x = 0; x < layer->width; x++) {
            BzTile ownerRawTile = data[y * layer->width + x];
            BzTileID owner = bzTilesetGetTileID(buildingTileset, ownerRawTile);
            BzTile buildingRawTile = buildingData[y * layer->width + x];
            BzTile buildingType = bzTilesetGetTileID(buildingTileset, buildingRawTile);
            buildingType = getTileBuilding(buildingType);
            if (buildingType <= BUILDING_NONE || buildingType >= BUILDING_COUNT)
                continue;
            // TODO: set player owner
            // Convert owner
            OwnerType ownerType = getOwnerType(owner);
            Player player = PLAYER_RED;
            switch (ownerType) {
                case OWNER_BLUE:
                    player = PLAYER_BLUE;
                    break;
                case OWNER_RED:
                    player = PLAYER_RED;
                    break;
                default:
                    BZ_ASSERT(false);
                    break;
            }
            placeBuilding(game, buildingType, x, y, player);

            i32 sizeX = 1;
            i32 sizeY = 1;
            getBuildingSize(buildingType, &sizeX, &sizeY);
            x += sizeX - 1;
            y += sizeY - 1;
        }
    }
    return true;

}

bool initRocksLayer(BzTileMap *map, BzTileLayer *layer) {
    Game *game = ecs_singleton_get_mut(ECS, Game);
    BzTileset *tileset = bzTileLayerGetTileset(map, layer);

    for (i32 y = 0; y < layer->height; y++) {
        for (i32 x = 0; x < layer->width; x++) {
            BzTile tile = layer->data[y * layer->width + x];
            BzTile layerTile = tile;
            tile = bzTilesetGetTileID(tileset, tile);
            if (tile == -1) continue; // Not a tree

            BzTileID tileID = bzTilesetGetTileID(tileset, layerTile);
            BZ_ASSERT(hasEntityHitBoxRec(tileID));
            HitBox hb = getEntityHitBoxRec(tileID);

            f32 sizeX = tileset->tileWidth;
            f32 sizeY = tileset->tileHeight;
            f32 posX = layer->offsetX + x * sizeX;
            f32 posY = layer->offsetY + y * sizeY;
            posY += sizeY;
            ecs_entity_t e = entityCreateEmpty();
            Position pos = {posX, posY};
            HitBox tHb = entityTransformHitBox(pos, hb);
            SpatialGridID gridID = bzSpatialGridInsert(game->entityGrid, &e,
                                                       tHb.x, tHb.y,
                                                       tHb.width, tHb.height);
            ecs_set(ECS, e, SpatialGridID, {gridID});
            ecs_set_ptr(ECS, e, Position, &pos);
            ecs_set(ECS, e, Size, {sizeX, sizeY});
            ecs_set_ptr(ECS, e, HitBox, &hb);
            ecs_set(ECS, e, Rotation, {0});
            ecs_set(ECS, e, TextureRegion, {tileset->tiles, getTextureRect(tileID)});
            ecs_set(ECS, e, Resource, {RES_GOLD, 80});
            ecs_add_id(ECS, e, Selectable);
            ecs_set(ECS, e, Harvestable, {
                    .harvestLimit = 4,
            });
        }
    }
    return true;
}

bool initTreesLayer(BzTileMap *map, BzTileLayer *layer) {
    Game *game = ecs_singleton_get_mut(ECS, Game);

    BzTileset *tileset = bzTileLayerGetTileset(map, layer);

    for (i32 y = 0; y < layer->height; y++) {
        for (i32 x = 0; x < layer->width; x++) {
            BzTile tile = layer->data[y * layer->width + x];
            BzTile layerTile = tile;
            tile = bzTilesetGetTileID(tileset, tile);
            if (tile == -1) continue; // Not a tree

            BzTile tileID = bzTilesetGetTileID(tileset, layerTile);
            BZ_ASSERT(hasEntityHitBoxRec(tileID));
            HitBox hb = getEntityHitBoxRec(tileID);

            f32 sizeX = tileset->tileWidth;
            f32 sizeY = tileset->tileHeight;
            f32 posX = layer->offsetX + x * sizeX;
            f32 posY = layer->offsetY + y * sizeY;
            posY += sizeY;
            Position pos = {posX, posY};
            ecs_entity_t e = entityCreateEmpty();
            Rectangle tHb = entityTransformHitBox((Position) {posX, posY}, hb);
            SpatialGridID gridID = bzSpatialGridInsert(game->entityGrid, &e,
                                                       tHb.x, tHb.y,
                                                       tHb.width, tHb.height);
            ecs_set(ECS, e, SpatialGridID, {gridID});
            ecs_set_ptr(ECS, e, Position, &pos);
            ecs_set(ECS, e, Size, {sizeX, sizeY});
            ecs_set_ptr(ECS, e, HitBox, &hb);
            ecs_set(ECS, e, Rotation, {0});
            ecs_set(ECS, e, TextureRegion, {tileset->tiles, getTextureRect(tileID)});
            ecs_set(ECS, e, Resource, {RES_WOOD, 20});
            ecs_add_id(ECS, e, Selectable);
            ecs_set(ECS, e, Harvestable, {
                    .harvestLimit = 4
            });

        }
    }
    return true;
}

void terrainRender(BzTileMap *map, BzTileLayer *layer) {
    BzTileset *tileset = bzTileLayerGetTileset(map, layer);


    Vector2 drawPos = {layer->offsetX, layer->offsetY};

    static f32 elapsed = 0.0f;
    elapsed += GetFrameTime();

    const Game *game = ecs_singleton_get(ECS, Game);
    Camera2D camera = game->camera;
    Rectangle camBounds = getCameraBounds(camera);

    for (i32 y = 0; y < layer->height; y++) {
        for (i32 x = 0; x < layer->width; x++) {
            BzTile tile = bzTileLayerGetTile(layer, x, y);
            tile        = bzTilesetGetTileID(tileset, tile);
            if (tile != -1) {
                if (terrainHasAnimation(tile)) {
                    f32 frameDuration = terrainGetAnimationFrame(tile, 0).duration;
                    i32 numFrames     = terrainGetAnimationSequence(tile).frameCount;
                    i32 frameIdx = (i32) (elapsed / frameDuration) % numFrames;
                    tile = terrainGetAnimationFrame(tile, frameIdx).frame;
                }
                Rectangle rec = bzTilesetGetTileRegion(tileset, tile);
                Rectangle bounds = {
                        drawPos.x,
                        drawPos.y,
                        tileset->tileWidth,
                        tileset->tileHeight
                };
                if (CheckCollisionRecs(camBounds, bounds))
                    DrawTextureRec(tileset->tiles, rec, drawPos, WHITE);
            }
            drawPos.x += (float) tileset->tileWidth;
        }
        drawPos.x = layer->offsetX;
        drawPos.y += (float) tileset->tileHeight;
    }
}
void loadMap(Game *game, const char *path, bool mainMenu) {

    game->map = bzTileMapCreate(&(BzTileMapDesc) {
            .path=path,
            .collisionMap=true,
            .tilesets[0]=game->tileset,

            .layers[LAYER_TERRAIN]=(BzTileLayerDesc) {"terrain", .renderer=terrainRender, .applyColliders=true},
            .layers[LAYER_ROCKS]=(BzTileLayerDesc) {"rocks", BZ_TILE_LAYER_SKIP_RENDER},
            .layers[LAYER_ROCKS2]=(BzTileLayerDesc) {"rocks_s", BZ_TILE_LAYER_SKIP_RENDER},
            .layers[LAYER_TREES]=(BzTileLayerDesc) {"trees", BZ_TILE_LAYER_SKIP_RENDER},
            .layers[LAYER_TREES2]=(BzTileLayerDesc) {"trees_s", BZ_TILE_LAYER_SKIP_RENDER},
            .layers[LAYER_BUILDINGS]=(BzTileLayerDesc) {"buildings", BZ_TILE_LAYER_SKIP_RENDER},
            .layers[LAYER_BUILDING_OWNER]=(BzTileLayerDesc) {"building_ownership", BZ_TILE_LAYER_SKIP_RENDER},

            .objectGroups[OBJECTS_GAME]=(BzTileObjectsDesc) {"game"},
            .objectGroups[OBJECTS_ENTITIES]=(BzTileObjectsDesc ) {"entities"}
    });
    game->entityGrid = bzSpatialGridCreate(&(BzSpatialGridDesc) {
            .maxWidth=game->map.width * game->map.tileWidth,
            .maxHeight=game->map.height * game->map.tileHeight,
            .cellWidth=game->map.tileWidth * 4,
            .cellHeight=game->map.tileHeight * 4,
            .userDataSize=sizeof(ecs_entity_t)
    });

    game->camera = (Camera2D){ 0 };
    game->camera.target = (Vector2) {0, 0};
    game->camera.offset = (Vector2) { GetScreenWidth() * 0.5f, GetScreenHeight() * 0.5f };
    game->camera.rotation = 0.0f;
    game->camera.zoom = 3.0f;

    bzMemSet(game->playerResources, 0, sizeof(*game->playerResources));
    game->waves = mainMenu ? getMainMenuWaves() : getDefaultWaves();
    game->waveInfo = getWaveInfo(&game->waves, 0);

    bzTileMapAddLayerCollisions(&game->map, LAYER_TERRAIN, COLL_LAYER_TERRAIN);

    bzTileMapOverrideLayer(&game->map, LAYER_TREES, initTreesLayer);
    bzTileMapOverrideLayer(&game->map, LAYER_TREES2, initTreesLayer);

    bzTileMapOverrideLayer(&game->map, LAYER_ROCKS, initRocksLayer);
    bzTileMapOverrideLayer(&game->map, LAYER_ROCKS2, initRocksLayer);

    bzTileMapOverrideLayer(&game->map, LAYER_BUILDING_OWNER, initBuildingsLayer);
    bzTileMapOverrideLayer(&game->map, LAYER_BUILDINGS, BZ_TILE_LAYER_CLEAR);

    bzTileMapOverrideObjectGroup(&game->map, OBJECTS_GAME, initGameObjectsLayer);
    bzTileMapOverrideObjectGroup(&game->map, OBJECTS_ENTITIES, initEntityObjectsLayer);


    if (mainMenu) {
        // Auto assign jobs
        i32 numWorkers = ecs_count_id(ECS, ecs_id(Worker));
        i32 numHarvestables = ecs_count_id(ECS, ecs_id(Harvestable));

        typedef struct PosPair {
            ecs_entity_t entity;
            Position pos;
        } PosPair;

        PosPair *workers = bzStackAlloc(&game->stackAlloc, sizeof(*workers) * numWorkers);
        PosPair *harvestables = bzStackAlloc(&game->stackAlloc, sizeof(*harvestables) * numWorkers);

        i32 workerIdx = 0;
        ecs_filter_t *workerFilter = ecs_filter(ECS, {
            .terms = {{ecs_id(Position)}, {ecs_id(Worker)}}
        });
        ecs_iter_t it = ecs_filter_iter(ECS, workerFilter);
        while (ecs_filter_next(&it)) {
            Position *pos = ecs_field(&it, Position, 1);
            for (i32 i = 0; i < it.count; i++) {
                if (workerIdx >= numWorkers)
                    break;
                workers[workerIdx++] = (PosPair) {
                    it.entities[i],
                    pos[i]
                };
            }
        }
        ecs_filter_fini(workerFilter);

        i32 harvestableIdx = 0;
        ecs_filter_t *harvestableFilter = ecs_filter(ECS, {
            .terms = {{ecs_id(Position)}, {ecs_id(Harvestable)}}
        });
        it = ecs_filter_iter(ECS, harvestableFilter);
        while (ecs_filter_next(&it)) {
            Position *pos = ecs_field(&it, Position, 1);
            for (i32 i = 0; i < it.count; i++) {
                if (harvestableIdx >= numHarvestables)
                    break;
                harvestables[harvestableIdx++] = (PosPair) {
                    it.entities[i],
                    pos[i]
                };
            }
        }
        ecs_filter_fini(harvestableFilter);
        bzLogInfo("%d %d", workerIdx, harvestableIdx);

        for (i32 i = 0; i < workerIdx; i++) {
            PosPair nearest = harvestables[0];
            f32 dst = INFINITY;
            for (i32 j = 1; j < harvestableIdx; j++) {
                if (nearest.entity == 0) continue;
                f32 jDst = Vector2Distance(workers[i].pos, harvestables[j].pos);
                if (jDst < dst) {
                    nearest = harvestables[j];
                    dst = jDst;
                }
            }

            if (nearest.entity == 0) continue;

            ResourceType resType = ecs_get(ECS, nearest.entity, Resource)->type;
            setAIBehaviour(workers[i].entity, game->BTs.workerHarvest, &(AIBlackboard) {
                .as.worker = {
                        .harvestType = resType,
                        .harvestTarget = nearest.entity,
                        .harvestPos = nearest.pos
                },
                .proximity = 8.0f,
            });
            Worker *worker = ecs_get_mut(ECS, workers[i].entity, Worker);
            worker->carryRes = resType;
        }

        bzStackAllocFree(&game->stackAlloc, harvestables);
        bzStackAllocFree(&game->stackAlloc, workers);


    }

}
void unloadMap(Game *game) {
    ecs_delete_with(ECS, GameEntity);
    if (game->map.isValid) {
        bzTileMapDestroy(&game->map);
        game->map.isValid = false;
    }
    if (game->entityGrid) {
        bzSpatialGridDestroy(game->entityGrid);
        game->entityGrid = NULL;
    }
}