#include <rlImGui.h>
#include <raygui.h>

#include "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 "unit_ai.h"
#include "unit_actions.h"

#include "pathfinding.h"

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

ECS_COMPONENT_DECLARE(Game);
ECS_COMPONENT_DECLARE(InputState);

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

static ecs_entity_t renderCollidersSystem;
static ecs_entity_t renderDebugPathSystem;

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) {
    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"},
            .layers[LAYER_TREES2]=(BzTileLayerDesc) {"trees_s"},
            .layers[LAYER_BUILDINGS]=(BzTileLayerDesc) {"buildings", .applyColliders=true},
            .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.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);

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

}

bool init(void *userData) {
    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);

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

    game->stackAlloc = bzStackAllocCreate(10 * 1000 * 1000); // 10 MB
    // init pools
    game->pools.pathData = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize = sizeof(PathData),
        .objectsPerPage = 512
    });
    game->pools.actions = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize = sizeof(Action),
        .objectsPerPage = 1024,
    });


    game->frameDuration = 0.16f;

    game->tileset = bzTilesetCreate( &(BzTilesetDesc) {
        .path="assets/game.tsj",
        .texturePath="assets/game.png"
    });
    loadMap(game, "assets/maps/main_menu_01.tmj");


    ECS_OBSERVER(ECS, entitySpatialRemove, EcsOnRemove, SpatialGridID);
    ECS_OBSERVER(ECS, entityPathRemove, EcsOnRemove, Path);

    //ECS_OBSERVER(ECS, entitySetAnimationState, EcsOnSet, Animation, AnimationType);

    ECS_SYSTEM(ECS, entityUpdateSpatialID, EcsOnUpdate, Position, Size, Velocity, SpatialGridID);
    ECS_SYSTEM(ECS, entityUpdateKinematic, EcsOnUpdate, Position, Rotation, Velocity, Steering);

    ECS_SYSTEM(ECS, entityMoveToTarget, EcsOnUpdate, Position, Rotation, Velocity, TargetPosition, Steering);
    ECS_SYSTEM(ECS, entityFollowPath, EcsOnUpdate, Path);

    //ECS_SYSTEM(ECS, entityHarvestTaskSystem, EcsOnUpdate, Position, Rotation, HarvestTask);
    ECS_SYSTEM(ECS, handleUnitActionsSystem, EcsOnUpdate, UnitAction);
    ECS_SYSTEM(ECS, updateUnitAISystem, EcsOnUpdate, UnitAI, UnitAction);
    // Needs to be called after AI update, since it removes finished actions
    ECS_SYSTEM(ECS, updateUnitActionsSystem, EcsOnUpdate, UnitAction);

    //ECS_SYSTEM(ECS, entityUpdateAnimationState, EcsOnUpdate, Velocity, AnimationType);
    ECS_SYSTEM(ECS, entityUpdateAnimation, EcsOnUpdate, Animation, TextureRegion);

    ECS_SYSTEM(ECS, renderDebugPath, EcsOnUpdate, Path);

    ECS_SYSTEM(ECS, renderTerrain, EcsOnUpdate, Position, Size, Rotation, TextureRegion, TextureTerrain);
    ECS_SYSTEM(ECS, renderBuildings, EcsOnUpdate, Position, Size, Rotation, TextureRegion, TextureBuildings);
    ECS_SYSTEM(ECS, renderEntities, EcsOnUpdate, Position, Size, Rotation, TextureRegion, TextureEntities);

    ECS_SYSTEM(ECS, renderColliders, EcsOnUpdate, Position, Size);
    ECS_SYSTEM(ECS, renderRotationDirection, EcsOnUpdate, Position, Rotation, TextureEntities);

    renderDebugPathSystem = renderDebugPath;
    renderCollidersSystem = renderColliders;

    game->debugDraw.mapColliders = true;
    game->debugDraw.spatialGrid = true;
    game->debugDraw.path = true;

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

    unloadMap(game);
    bzTilesetDestroy(&game->tileset);

    Game gameCopy = *game;
    InputState inputCopy = *input;

    // Destroy queries
    ecs_query_fini(inputCopy.queries.selected);

    ecs_fini(ECS);
    ECS = NULL;

    bzStackAllocDestroy(&gameCopy.stackAlloc);
    bzObjectPoolDestroy(gameCopy.pools.pathData);
    bzObjectPoolDestroy(gameCopy.pools.actions);

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

    updateInputState(input, game->camera, dt);

    switch (game->screen) {
        case SCREEN_GAME:
            updatePlayerInput();
            break;
        case SCREEN_MAIN_MENU:
            break;
        case SCREEN_SETTINGS:
            break;
    }
}

static void renderGame(Game *game, float dt) {
    ClearBackground(RAYWHITE);
    BeginMode2D(game->camera);

    bzTileMapDraw(&game->map);

    drawPlayerInputUIGround();

    ecs_progress(ECS, dt);
    ecs_enable(ECS, renderDebugPathSystem, game->debugDraw.path);
    ecs_enable(ECS, renderCollidersSystem, game->debugDraw.entityColliders);
    if (game->debugDraw.mapColliders)
        bzTileMapDrawCollisions(&game->map);
    if (game->debugDraw.spatialGrid)
        bzSpatialGridDrawDebugGrid(game->entityGrid);

    drawPlayerInputUI();

    EndMode2D();

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

    DrawRectangle(0, 0, width, height, c);
}
static void renderMainMenu(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("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")) {
        bzLogInfo("Play");
        game->screen = SCREEN_GAME;
    }
    if (uiMainMenuButton("Settings")) {
        bzLogInfo("Settings");
    }
    if (uiMainMenuButton("Exit")) {
        bzLogInfo("Bye");
        bzGameExit();
    }
    bzUIPopParent(UI);
    bzUIEnd(UI);

}

static void renderSettings(Game *game, float dt) {
}

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

    Color shadow = BLACK;
    shadow.a = 35;

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


}

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

    if (game->screen != SCREEN_GAME) return;

    igSetNextWindowSize((ImVec2){300, 400}, ImGuiCond_FirstUseEver);
    igBegin("Debug Menu", NULL, 0);
    if (igSmallButton("Recruit worker [50 food]")) {
        createWorker((Position) {1100, 400}, (Size) {10, 10}, game->entityGrid,
                     &game->map.tilesets[2], 1322);
    }
    igText("PathData pool available: %llu", bzObjectPoolGetNumFree(game->pools.pathData));
    igText("Action   pool available: %llu", bzObjectPoolGetNumFree(game->pools.actions));
    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: {
                igText("Selected units:");
                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 %llu", entity);
                        igText("Actions:");
                        const Action *pAction = NULL;
                        if (ecs_has(ECS, entity, UnitAction)) {
                            pAction = ecs_get(ECS, entity, UnitAction)->first;
                        }
                        while (pAction) {
                            igText("\t%d: %s", pAction->type, actionTypeToPrettyStr(pAction->type));
                            pAction = pAction->next;
                        }
                    }
                }
                break;
            }
            case INPUT_SELECTED_OBJECT: {
                igText("Selected objects:");
                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];
                        if (ecs_has(ECS, entity, Harvestable) &&
                            ecs_has(ECS, entity, Resource)) {
                            Resource *res = ecs_get_mut(ECS, entity, Resource);
                            const char *resName = getResourceTypePrettyName(res->type);
                            igText("\tEntity %llu:", entity);
                            igSliderInt("\t\twood", &res->amount, 0, 20, NULL, 0);
                        }
                    }
                }
                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 = 0; i < BUILDING_COUNT; i++) {
            if (igSelectable_Bool(getBuildingStr(i), input->building == i, 0, (ImVec2){0, 0}))
                input->building = i;
        }
        if (input->building)
            input->state = INPUT_BUILDING;

    }
    if (igCollapsingHeader_TreeNodeFlags("DebugDraw", 0)) {
        igCheckbox("map colliders", &game->debugDraw.mapColliders);
        igCheckbox("entity colliders", &game->debugDraw.entityColliders);
        igCheckbox("spatial grid", &game->debugDraw.spatialGrid);
        igCheckbox("path", &game->debugDraw.path);

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