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

#include "systems/systems.h"
#include "components.h"
#include "game_state.h"
#include "game_tileset.h"
#include "input.h"
#include "map_init.h"
#include "map_layers.h"
#include "buildings.h"
#include "ui_widgets.h"

#include "pathfinding.h"
#include "sounds.h"


ECS_COMPONENT_DECLARE(Game);
ECS_COMPONENT_DECLARE(InputState);
ECS_COMPONENT_DECLARE(SoundState);

BzUI *UI = NULL;
ecs_world_t *ECS = NULL;

bool init(void *userData);
void deinit(void *userData);

void update(float dt, void *userData);
void render(float dt, void *userData);
void imguiRender(float dt, void *userData);


bool bzMain(BzAppDesc *appDesc, int argc, const char **argv) {
    appDesc->width = 1280;
    appDesc->height = 720;
    appDesc->title = "PixelDefense";
    appDesc->fps = 60;

    appDesc->init = (BzAppInitFunc) init;
    appDesc->deinit = (BzAppDeinitFunc) deinit;
    appDesc->update = (BzAppUpdateFunc) update;
    appDesc->render = (BzAppRenderFunc) render;
    appDesc->imguiRender = (BzAppRenderFunc) imguiRender;

    appDesc->userData = NULL;

    //SetConfigFlags(FLAG_WINDOW_RESIZABLE);

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

    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);
                DrawTextureRec(tileset->tiles, rec, drawPos, WHITE);
            }
            drawPos.x += (float) tileset->tileWidth;
        }
        drawPos.x = layer->offsetX;
        drawPos.y += (float) tileset->tileHeight;
    }
}

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;
    }
}
void loadMap(Game *game, const char *path) {
    game->map = bzTileMapCreate(&(BzTileMapDesc) {
            .path=path,
            .generateCollisionMap=true,
            .tilesets[0]=game->tileset,

            .layers[LAYER_TERRAIN]=(BzTileLayerDesc) {"terrain", .renderer=terrainRender, .applyColliders=true},
            .layers[LAYER_ROCKS]=(BzTileLayerDesc) {"rocks"},
            .layers[LAYER_ROCKS2]=(BzTileLayerDesc) {"rocks_s"},
            .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;

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

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

}
int cmpDrawData(const void *a, const void *b) {
    const DrawData *lhs = (DrawData *) a;
    const DrawData *rhs = (DrawData *) b;

    f32 dif = (rhs->dst.y) - (lhs->dst.y);
    int cmpVal = 0;
    if (dif < 0) cmpVal = 1;
    else if (dif > 0) cmpVal = -1;
    return cmpVal;
}

bool init(void *userData) {
    // Center window
    int monitor = GetCurrentMonitor();
    int monitorWidth = GetMonitorWidth(monitor);
    int monitorHeight = GetMonitorHeight(monitor);
    int windowWidth = GetScreenWidth();
    int windowHeight = GetScreenHeight();
    SetWindowPosition((int)(monitorWidth / 2) - (int)(windowWidth / 2), (int)(monitorHeight / 2) - (int)(windowHeight / 2));

    BZ_UNUSED(userData);
    SetExitKey(0);

    UI = bzUICreate();

    ECS = ecs_init();

    initComponentIDs(ECS);

    // For ecs explorer
    //ecs_singleton_set(ECS, EcsRest, {0});
    //ECS_IMPORT(ECS, FlecsMonitor);

    ECS_COMPONENT_DEFINE(ECS, Game);
    ecs_singleton_set(ECS, Game, {});
    Game *game = ecs_singleton_get_mut(ECS, Game);
    game->screen = SCREEN_MAIN_MENU;
    game->font = LoadFontEx("assets/fonts/CompassPro.ttf", 92, NULL, 0);
    game->drawData = bzArrayCreate(DrawData, 1000);
    game->drawQuery = ecs_query(ECS, {
        .filter.terms = {
                { ecs_id(Position) },
                { ecs_id(Size) },
                { ecs_id(Rotation) },
                { ecs_id(TextureRegion) }
        },
    });

    {
        ECS_COMPONENT_DEFINE(ECS, InputState);
        ecs_singleton_set(ECS, InputState, {});
        InputState *input = ecs_singleton_get_mut(ECS, InputState);

        input->mapping = inputDefaultMapping();

        // Create queries
        input->queries.selected = ecs_query(ECS, {
            .filter.terms = {
                    {ecs_id(Position)},
                    {ecs_id(Size)},
                    {ecs_id(Selected)}
            }
        });

        input->building = BUILDING_NONE;
    }
    {
        InitAudioDevice();
        ECS_COMPONENT_DEFINE(ECS, SoundState);
        ecs_singleton_set(ECS, SoundState, {});
        SoundState *sounds = ecs_singleton_get_mut(ECS, SoundState);
        soundsLoad(sounds, SOUND_WOOD_PUNCH, 0.1f, "assets/sounds/wood hit 17.wav");
    }

    game->stackAlloc = bzStackAllocCreate(10 * 1000 * 1000); // 10 MB
    // init pools
    game->pools.pathData = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize = sizeof(PathData),
        .objectsPerPage = 512
    });
    game->pools.btNode = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize = bzBTGetNodeSize(),
        .objectsPerPage = 64
    });
    game->pools.btNodeState = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize = bzBTGetNodeStateSize(),
        .objectsPerPage = 1024,
    });
    BzObjectPool *nodePool = game->pools.btNode;
    // moveTo
    {
        BzBTNode *root = NULL;
        BzBTNode *node = NULL;
        root = bzBTMakeRoot(nodePool);
        game->BTs.moveTo = root;

        // Just a single action for now
        node = bzBTAction(nodePool, root, (BzBTActionFn) aiMoveTo);
        bzBTNodeSetName(node, "moveTo");
    }
    // worker harvest
    {
        BzBTNode *root = NULL;
        BzBTNode *node = NULL;
        root = bzBTMakeRoot(nodePool);
        game->BTs.workerHarvest = root;

        //node = bzBTDecorUntilFail(nodePool, root);

        BzBTNode *collectSeq = bzBTCompSequence(nodePool, root, false);

        BzBTNode *untilFail = bzBTDecorUntilFail(nodePool, collectSeq);
        {
            BzBTNode *untilSeq = bzBTCompSequence(nodePool, untilFail, false);
            node = bzBTAction(nodePool, untilSeq, (BzBTActionFn) aiFindNextHarvestable);
            bzBTNodeSetName(node, "findNextHarvestable");
            node = bzBTAction(nodePool, untilSeq, (BzBTActionFn) aiMoveTo);
            bzBTNodeSetName(node, "moveTo");
            node = bzBTAction(nodePool, untilSeq, (BzBTActionFn) aiResetElapsed);
            bzBTNodeSetName(node, "resetElapsed");
            node = bzBTAction(nodePool, untilSeq, (BzBTActionFn) aiHarvestRes);
            bzBTNodeSetName(node, "harvestRes");
            node = bzBTDecorInvert(nodePool, untilSeq);
            node = bzBTAction(nodePool, node, (BzBTActionFn) aiCarryCapacityFull);
            bzBTNodeSetName(node, "carryCapacityFull");
        }
        node = bzBTDecorInvert(nodePool, collectSeq);
        node = bzBTAction(nodePool, node, (BzBTActionFn) aiCarryCapacityEmpty);
        bzBTNodeSetName(node, "carryCapacityEmpty");

        node = bzBTAction(nodePool, collectSeq, (BzBTActionFn) aiFindNearestStorage);
        bzBTNodeSetName(node, "findNearestStorage");
        node = bzBTAction(nodePool, collectSeq, (BzBTActionFn) aiMoveTo);
        bzBTNodeSetName(node, "moveTo");
        node = bzBTAction(nodePool, collectSeq, (BzBTActionFn) aiResetElapsed);
        bzBTNodeSetName(node, "resetElapsed");
        node = bzBTAction(nodePool, collectSeq, (BzBTActionFn) aiDepositRes);
        bzBTNodeSetName(node, "depositRes");
        node = bzBTDecorDelay(nodePool, collectSeq, 1.0f);


        //node = bzBTAction(nodePool, collectSeq, NULL);
        //bzBTNodeSetName(node, "harvest");
        //node = bzBTAction(nodePool, collectSeq, NULL);
        //bzBTNodeSetName(node, "moveTo");
        //node = bzBTAction(nodePool, collectSeq, NULL);
        //bzBTNodeSetName(node, "deposit");
    }


    game->frameDuration = 0.16f;

    game->tileset = bzTilesetCreate( &(BzTilesetDesc) {
        .path="assets/game.tsj",
        .texturePath="assets/game.png"
    });
    // Fixes texture bleeding
    GenTextureMipmaps(&game->tileset.tiles);
    SetTextureWrap(game->tileset.tiles, TEXTURE_WRAP_CLAMP);
    SetTextureFilter(game->tileset.tiles, TEXTURE_FILTER_POINT);

    setupSystems();
    loadMap(game, "assets/maps/main_menu_01.tmj");


    game->debug.drawMapColliders = true;
    game->debug.drawSpatialGrid = true;
    game->debug.drawPath = true;
    game->debug.inspecting = bzArrayCreate(ecs_entity_t, 10);

    return true;
}
void deinit(void *userData) {
    BZ_UNUSED(userData);
    Game *game = ecs_singleton_get_mut(ECS, Game);
    InputState *input = ecs_singleton_get_mut(ECS, InputState);
    SoundState *sounds = ecs_singleton_get_mut(ECS, SoundState);

    unloadMap(game);

    // Destroy queries
    ecs_query_fini(input->queries.selected);
    ecs_query_fini(game->drawQuery);

    Game gameCopy = *game;
    InputState inputCopy = *input;
    SoundState soundsCopy = *sounds;

    // Destroy ECS
    ecs_fini(ECS);
    ECS = NULL;

    game = &gameCopy;
    input = &inputCopy;
    sounds = &soundsCopy;

    bzArrayDestroy(game->debug.inspecting);

    bzTilesetDestroy(&game->tileset);

    bzStackAllocDestroy(&game->stackAlloc);
    bzObjectPoolDestroy(game->pools.pathData);
    bzObjectPoolDestroy(game->pools.btNode);
    bzObjectPoolDestroy(game->pools.btNodeState);

    bzArrayDestroy(game->drawData);

    soundsUnloadAll(sounds);

    bzUIDestroy(UI);
    UI = NULL;

    UnloadFont(game->font);
}


void update(float dt, void *userData) {
    BZ_UNUSED(userData);

    char titleBuf[32];
    snprintf(titleBuf, sizeof(titleBuf), "FPS: %d | %.2f ms", GetFPS(), GetFrameTime() * 1000);
    SetWindowTitle(titleBuf);

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

    BZ_ASSERT(game->stackAlloc.allocated == 0);
    bzStackAllocReset(&game->stackAlloc);

    ImGuiIO *io = igGetIO();
    bzUISetCaptureMouse(UI, !io->WantCaptureMouse);
    bzUISetCaptureKeyboard(UI, !io->WantCaptureKeyboard);

    input->canUseMouse = !io->WantCaptureMouse && !bzUICapturedMouse(UI);
    input->canUseKeyboard = !io->WantCaptureKeyboard && !bzUICapturedKeyboard(UI);
    updateInputState(input, game->camera, dt);

    switch (game->screen) {
        case SCREEN_GAME:
            updatePlayerInput();
            break;
        case SCREEN_PAUSE_MENU:
            if (IsKeyReleased(input->mapping.backBtn)) {
                game->screen = SCREEN_GAME;
            }
            break;
        case SCREEN_MAIN_MENU:
            break;
        case SCREEN_SETTINGS:
            break;
    }
}

static void drawOverScreen(Color c) {
    i32 width = GetScreenWidth();
    i32 height = GetScreenHeight();

    DrawRectangle(0, 0, width, height, c);
}
static void renderGame(Game *game, float dt) {
    ClearBackground(RAYWHITE);
    BeginMode2D(game->camera);

    // Map
    bzTileMapDraw(&game->map);
    // Ground UI
    drawPlayerInputUIGround();
    // Entities
    bzArrayClear(game->drawData);
    ecs_iter_t it = ecs_query_iter(ECS, game->drawQuery);
    ecs_entity_t worker = 0;
    while (ecs_iter_next(&it)) {
        Position *p = ecs_field(&it, Position, 1);
        Size *s = ecs_field(&it, Size, 2);
        Rotation *r = ecs_field(&it, Rotation, 3);
        TextureRegion *t = ecs_field(&it, TextureRegion, 4);
        for (i32 i = 0; i < it.count; i++) {
            Rectangle dst = {p[i].x, p[i].y, s[i].x, s[i].y};
            if (dst.width == 10 && dst.height == 10) {
                worker = it.entities[i];
            }
            Vector2 origin = {dst.width * 0.5f, dst.height};
            dst.x += origin.x - dst.width * 0.5f;
            dst.y += origin.y - dst.height * 0.5f;
            Rectangle src = t[i].rec;
            // Fixes texture bleeding issue
            src.x += 0.01f;
            src.y += 0.01f;
            src.width -= 0.02f;
            src.height -= 0.02f;
            if (t[i].flipX) src.width *= -1.0f;
            if (t[i].flipY) src.height *= -1.0f;
            bzArrayPush(game->drawData, (DrawData) {
                .tex = t[i].texture,
                .src = src,
                .dst = dst,
                .origin = origin,
                .rotation = r[i]
            });
        }
    }
    qsort(game->drawData, bzArraySize(game->drawData), sizeof(*game->drawData), cmpDrawData);

    for (i32 i = 0; i < bzArraySize(game->drawData); i++) {
        DrawData draw = game->drawData[i];
        DrawTexturePro(draw.tex, draw.src, draw.dst, draw.origin, draw.rotation, WHITE);
    }
    Vector2 target = GetMousePosition();
    target = GetScreenToWorld2D(target, game->camera);
    static f32 elapsed = 0;
    static bool attack = false;
    static Vector2 lockedTarget;
    if (!attack && IsMouseButtonPressed(0)) {
        attack = true;
        lockedTarget = target;
        elapsed = 0;
    }
    elapsed += dt * 2;
    elapsed = Clamp(elapsed, 0, 1.0f);
    attack = false;
    if (worker && false) {
        Position *pos = ecs_get_mut(ECS, worker, Position);
        DrawCircle(pos->x, pos->y, 2.0f, BLUE);
        Vector2 attackVector = Vector2Subtract(lockedTarget, *pos);
        attackVector = Vector2Normalize(attackVector);
        attackVector = Vector2Scale(attackVector, 2.0f);
        DrawLine(pos->x, pos->y, pos->x + attackVector.x, pos->y + attackVector.y, RED);
        Rotation *rot = ecs_get_mut(ECS, worker, Rotation);
        f32 targetRot = Vector2Angle(*pos, lockedTarget);
        targetRot += 25 * DEG2RAD;
        *rot = targetRot * bzEase(BZ_EASE_IN_BACK, elapsed);
        bzLogInfo("%.2f", Vector2Angle(*pos, lockedTarget) * RAD2DEG);
    }

    ecs_progress(ECS, dt);
    ecs_enable(ECS, renderDebugPathSystem, game->debug.drawPath);
    ecs_enable(ECS, renderCollidersSystem, game->debug.drawEntityColliders);
    if (game->debug.drawMapColliders)
        bzTileMapDrawCollisions(&game->map);
    if (game->debug.drawSpatialGrid)
        bzSpatialGridDrawDebugGrid(game->entityGrid);
    drawPlayerInputUI();

    EndMode2D();
}

static void renderGameMenu(Game *game, float dt) {
    // UI
    i32 width = GetScreenWidth();
    i32 height = GetScreenHeight();
    bzUIBegin(UI, width, height);

    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN |
                     BZ_UI_FLEX_JUSTIFY_SPACE_BETWEEN |
                     BZ_UI_FLEX_ALIGN_CENTER
    });
    // top bar
    f32 topBarHeight = 0.05f;
    BzUINode *topBar = uiPushDivParentPercentage(1.0f, topBarHeight);
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_ROW |
                     BZ_UI_FLEX_ALIGN_CENTER |
                     BZ_UI_FLEX_JUSTIFY_START
    });
    Color topBarBG = {0, 0, 0, 50};
    bzUISetBackgroundStyle(UI, topBar, (BzUIBackgroundStyle) {
            .normal = topBarBG,
            .hover = topBarBG,
            .active = topBarBG,
    });
    BzTileset *tileset = &game->tileset;
    Rectangle woodRec  = bzTilesetGetTileRegion(tileset, getEntityTile(ENTITY_WOOD));
    Rectangle stoneRec = bzTilesetGetTileRegion(tileset, getEntityTile(ENTITY_STONE));
    Rectangle foodRec  = bzTilesetGetTileRegion(tileset, getEntityTile(ENTITY_APPLE));
    Rectangle goldRec  = bzTilesetGetTileRegion(tileset, getEntityTile(ENTITY_GOLD));
    Rectangle popRec  = bzTilesetGetTileRegion(tileset, getEntityTile(ENTITY_POP));
    uiGameResCount(100, -1, woodRec, tileset->tiles);
    uiGameResCount(100, -1, stoneRec, tileset->tiles);
    uiGameResCount(100, -1, foodRec, tileset->tiles);
    uiGameResCount(250, -1, goldRec, tileset->tiles);
    uiGameResCount(1, 10, popRec, tileset->tiles);
    bzUIPopParent(UI);


    bzUIEnd(UI);
}
static void renderPauseMenu(Game *game, float dt) {
    i32 width = GetScreenWidth();
    i32 height = GetScreenHeight();

    bzUIBegin(UI, width, height);
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN
    });

    uiPushDivParentPercentage(1.0f, 0.4f);
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_JUSTIFY_CENTER | BZ_UI_FLEX_ALIGN_CENTER
    });
    uiMainMenuLabel("Paused");
    bzUIPopParent(UI);

    uiPushDivParentPercentage(1.0f, 0.6f);
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_ALIGN_CENTER
    });
    if (uiMainMenuButton("Resume")) {
        game->screen = SCREEN_GAME;
    }
    if (uiMainMenuButton("Exit")) {
        game->screen = SCREEN_MAIN_MENU;
        unloadMap(game);
        loadMap(game, "assets/maps/main_menu_01.tmj");
    }
    bzUIPopParent(UI);
    bzUIEnd(UI);
}

static void renderMainMenu(Game *game, float dt) {
    i32 width = GetScreenWidth();
    i32 height = GetScreenHeight();

    game->camera.zoom = 3 * uiGetScale();

    bzUIBegin(UI, width, height);
    bzUISetParentLayout(UI, (BzUILayout) {
        .type = BZ_UI_LAYOUT_FLEX_BOX,
        .flags = BZ_UI_FLEX_DIR_COLUMN
    });

    uiPushDivParentPercentage(1.0f, 0.4f);
    bzUISetParentLayout(UI, (BzUILayout) {
        .type = BZ_UI_LAYOUT_FLEX_BOX,
        .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_JUSTIFY_CENTER | BZ_UI_FLEX_ALIGN_CENTER
    });
    uiMainMenuLabel("Pixel Defense");
    bzUIPopParent(UI);

    uiPushDivParentPercentage(1.0f, 0.6f);
    bzUISetParentLayout(UI, (BzUILayout) {
        .type = BZ_UI_LAYOUT_FLEX_BOX,
        .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_ALIGN_CENTER
    });
    if (uiMainMenuButton("Play")) {
        game->screen = SCREEN_GAME;
        unloadMap(game);
        loadMap(game, "assets/maps/pathing_test.tmj");
    }
    if (uiMainMenuButton("Settings")) {
        game->screen = SCREEN_SETTINGS;
    }
    if (uiMainMenuButton("Exit")) {
        bzGameExit();
    }
    bzUIPopParent(UI);
    bzUIEnd(UI);

}

static void renderSettings(Game *game, float dt) {
    i32 width = GetScreenWidth();
    i32 height = GetScreenHeight();

    game->camera.zoom = 3 * uiGetScale();

    bzUIBegin(UI, width, height);
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_JUSTIFY_CENTER | BZ_UI_FLEX_ALIGN_CENTER
    });

    bzUIPushDiv(UI, (BzUISize) { BZ_UI_SIZE_REL_PARENT, 0.8f},
                (BzUISize) { BZ_UI_SIZE_REL_PARENT, 0.8f});
    bzUISetParentLayout(UI, (BzUILayout) {
            .type = BZ_UI_LAYOUT_FLEX_BOX,
            .flags = BZ_UI_FLEX_DIR_COLUMN | BZ_UI_FLEX_JUSTIFY_CENTER | BZ_UI_FLEX_ALIGN_CENTER
    });
    static bool fullscreen = false;
    static bool vsync = false;
    uiSettingsLabel("Video");
    uiSettingsCheckbox("Fullscreen", &fullscreen);
    uiSettingsCheckbox("V-Sync", &vsync);

    static f32 master = 50.0f;
    static f32 music  = 50.0f;
    static f32 sound  = 50.0f;
    uiSettingsLabel("Audio");
    uiSettingsSlider("Master:  ", &master);
    uiSettingsSlider("Music:   ", &music);
    uiSettingsSlider("Sound:   ", &sound);

    bzUIPopParent(UI);
    bzUIPushDiv(UI, (BzUISize) {BZ_UI_SIZE_REL_PARENT, 0.8f},
                (BzUISize) {BZ_UI_SIZE_REL_PARENT, 0.2f});
    bzUISetParentLayout(UI, (BzUILayout) {
        .type = BZ_UI_LAYOUT_FLEX_BOX,
        .flags = BZ_UI_FLEX_DIR_ROW | BZ_UI_FLEX_JUSTIFY_CENTER | BZ_UI_FLEX_ALIGN_CENTER
    });

    if (uiSettingsButton("Back")) {
        game->screen = SCREEN_MAIN_MENU;
    }
    if (uiSettingsButton("Reset")) {

    }
    if (uiSettingsButton("Apply")) {
        game->screen = SCREEN_MAIN_MENU;
    }

    bzUIEnd(UI);
}

void render(float dt, void *userData) {
    BZ_UNUSED(userData);
    Game *game = ecs_singleton_get_mut(ECS, Game);
    const InputState *input = ecs_singleton_get(ECS, InputState);;

    Color shadow = BLACK;
    shadow.a = 35;

    switch (game->screen) {
        case SCREEN_GAME:
            renderGame(game, dt);
            renderGameMenu(game, dt);
            break;
        case SCREEN_PAUSE_MENU:
            renderGame(game, dt);
            drawOverScreen(shadow);
            renderPauseMenu(game, dt);
            break;
        case SCREEN_MAIN_MENU:
            renderGame(game, dt);
            drawOverScreen(shadow);
            renderMainMenu(game, dt);
            break;
        case SCREEN_SETTINGS:
            renderGame(game, dt);
            drawOverScreen(shadow);
            renderSettings(game, dt);
            break;
    }


}

void igInspectComp(const char *label, ecs_entity_t entity, ecs_entity_t comp, ImGuiCompFn fn) {
    igPushID_Int(comp);
    if (igTreeNode_Str(label)) {
        bool isAttached = ecs_has_id(ECS, entity, comp);
        igCheckbox("Attached", &isAttached);

        if (isAttached)
            fn(ECS, entity, comp);

        if (isAttached != ecs_has_id(ECS, entity, comp)) {
            if (!isAttached) {
                ecs_remove_id(ECS, entity, comp);
            } else {
                ecs_set_id(ECS, entity, comp, 0, NULL);
            }
        }
        igTreePop();
    }
    igPopID();
}
void igInspectWindow(ecs_entity_t entity, bool *open) {
    if (!ecs_is_alive(ECS, entity)) {
        *open = false;
        return;
    }
    igSetNextWindowSize((ImVec2) {300, 440}, ImGuiCond_FirstUseEver);
    char buf[64];
    snprintf(buf, sizeof(buf), "Entity: %ld", entity);
    if (igBegin(buf, open, 0)) {
        if (igSmallButton("Destroy")) {
            ecs_delete(ECS, entity);
            igEnd();
            *open = false;
            return;
        }
        if (igCollapsingHeader_TreeNodeFlags("Tags", 0)) {
            //igTagCheckbox("GameEntity", ECS, entity, GameEntity);
            igTagCheckbox("Selectable", ECS, entity, Selectable);
            igTagCheckbox("Selected", ECS, entity, Selected);
            igTagCheckbox("Storage", ECS, entity, Storage);
            igTagCheckbox("Harvestable", ECS, entity, Harvestable);
            igTagCheckbox("Workable", ECS, entity, Workable);
            igTagCheckbox("Attackable", ECS, entity, Attackable);
        }
        if (ecs_has(ECS, entity, BzBTState) &&
            igCollapsingHeader_TreeNodeFlags("BehaviourTree", 0)) {
            const BzBTState *state = ecs_get(ECS, entity, BzBTState);
            if (state->root)
                igVisualizeBTState(state->root, state->_first, true, false, 0);
            else
                igTextColored((ImVec4) {1, 0, 0, 1}, "NONE");
        }
        if (igCollapsingHeader_TreeNodeFlags("Components", 0)) {
            igInspectComp("Resource", entity, ecs_id(Resource), igResource);
            igInspectComp("Owner", entity, ecs_id(Owner), igOwner);
            igInspectComp("SpatialGridID", entity, ecs_id(SpatialGridID), igSpatialGridID);
            igInspectComp("Position", entity, ecs_id(Position), igVec2Comp);
            igInspectComp("Size", entity, ecs_id(Size), igVec2Comp);
            igInspectComp("Velocity", entity, ecs_id(Velocity), igVec2Comp);
            igInspectComp("TargetPosition", entity, ecs_id(TargetPosition), igVec2Comp);
            igInspectComp("Steering", entity, ecs_id(Steering), igVec2Comp);
            igInspectComp("Rotation", entity, ecs_id(Rotation), igFloat);
            igInspectComp("Path", entity, ecs_id(Path), igPath);
            igInspectComp("TextureRegion", entity, ecs_id(TextureRegion), igTextureRegion);
            igInspectComp("Animation", entity, ecs_id(Animation), igAnimation);
            igInspectComp("Easing", entity, ecs_id(Easing), igEasing);
            igInspectComp("Arms", entity, ecs_id(Arms), igArms);
            igInspectComp("Arm", entity, ecs_id(Arm), igArm);
            igInspectComp("BzBTState", entity, ecs_id(BzBTState), igBzBTState);
            igInspectComp("AIBlackboard", entity, ecs_id(AIBlackboard), igAIBlackboard);
            igInspectComp("Worker", entity, ecs_id(Worker), igWorker);
            igInspectComp("Unit", entity, ecs_id(Unit), igUnit);
        }
    }
    igEnd();
}
void imguiRender(float dt, void *userData) {
    BZ_UNUSED(userData);
    Game *game = ecs_singleton_get_mut(ECS, Game);
    InputState *input = ecs_singleton_get_mut(ECS, InputState);

    igShowDemoWindow(NULL);

    igSetNextWindowSize((ImVec2){300, 400}, ImGuiCond_FirstUseEver);
    igBegin("Debug Menu", NULL, 0);
    igText("PathData    pool available: %llu", bzObjectPoolGetNumFree(game->pools.pathData));
    igText("BTNode      pool available: %llu", bzObjectPoolGetNumFree(game->pools.btNode));
    igText("BTNodeState pool available: %llu", bzObjectPoolGetNumFree(game->pools.btNodeState));
    const char *inputState = "NONE";
    switch (input->state) {
        case INPUT_NONE:
            inputState = "NONE";
            break;
        case INPUT_BUILDING:
            inputState = "BUILDING";
            break;
        case INPUT_SELECTED_UNITS:
            inputState = "SELECTED_UNITS";
            break;
        case INPUT_SELECTED_OBJECT:
            inputState = "SELECTED_OBJECT";
            break;
        case INPUT_SELECTED_BUILDING:
            inputState = "SELECTED_BUILDING";
            break;
    }
    igText("Input state: %s", inputState);
    if (igCollapsingHeader_TreeNodeFlags("Selection", 0)) {
        switch (input->state) {
            case INPUT_SELECTED_UNITS:
            case INPUT_SELECTED_OBJECT:
            case INPUT_SELECTED_BUILDING: {
                ecs_iter_t it = ecs_query_iter(ECS, input->queries.selected);
                while (ecs_iter_next(&it)) {
                    for (i32 i = 0; i < it.count; i++) {
                        ecs_entity_t entity = it.entities[i];
                        igText("Entity: %ld", entity);
                    }
                }
                break;
            }
            default:
                igText("NONE");
                break;
        }

    }
    if (igCollapsingHeader_TreeNodeFlags("Resources", 0)) {
        igText("Wood:       %lld", game->resources.wood);
        igText("Iron:       %lld", game->resources.iron);
        igText("Food:       %lld", game->resources.food);
        igText("Gold:       %lld", game->resources.gold);
        igText("Population: %lld", game->resources.pop);
    }
    if (igCollapsingHeader_TreeNodeFlags("BuildMenu", 0)) {
        for (int i = BUILDING_NONE; i < BUILDING_COUNT; i++) {
            const char *buildingStr = getBuildingStr(i);
            if (!buildingStr) buildingStr = "NONE";
            if (igSelectable_Bool(buildingStr, input->building == i, 0, (ImVec2){0, 0}))
                input->building = i;
        }
        if (input->building > BUILDING_NONE && input->building < BUILDING_COUNT)
            input->state = INPUT_BUILDING;

    }
    if (igCollapsingHeader_TreeNodeFlags("DebugDraw", 0)) {
        igCheckbox("map colliders", &game->debug.drawMapColliders);
        igCheckbox("entity colliders", &game->debug.drawEntityColliders);
        igCheckbox("spatial grid", &game->debug.drawSpatialGrid);
        igCheckbox("path", &game->debug.drawPath);

    }
    if (igCollapsingHeader_TreeNodeFlags("Entities", 0)) {
        igSliderFloat("Frame duration", &game->frameDuration, 0.0f, 1.0f, NULL, 0);
    }
    if (igSmallButton("Quite game")) {
        bzGameExit();
    }
    igEnd();

    ecs_defer_begin(ECS);
    i32 inspectLen = bzArraySize(game->debug.inspecting);
    for (i32 i = inspectLen - 1; i >= 0; i--) {
        bool open = true;
        igInspectWindow(game->debug.inspecting[i], &open);
        if (!open) {
            bzArrayDelSwap(game->debug.inspecting, i);
        }
    }
    ecs_defer_end(ECS);
}