#ifndef WORLDBOX_CPP
#define WORLDBOX_CPP

#include "../util/sim/worldbox.hpp"
#include "../util/grid/camera.hpp"

#include "../imgui/imgui.h"
#include "../imgui/backends/imgui_impl_glfw.h"
#include "../imgui/backends/imgui_impl_opengl3.h"
#include <GLFW/glfw3.h>

class worldboxSimUI {
private:
    worldboxsim sim;
    Camera cam;
    
    // UI and Render State
    GLuint textu = 0;
    std::mutex PreviewMutex;
    bool updatePreview = false;
    bool textureInitialized = false;
    frame currentPreviewFrame;
    
    // Rendering Settings
    int outWidth = 1024;
    int outHeight = 1024;
    int reflectCount = 2;
    bool slowRender = false;
    float lodDist = 1024.0f;
    float lodDropoff = 0.05f;
    float maxViewDistance = 4096.0f;
    bool globalIllumination = false;
    bool useLod = true;
    float framerate = 60.0f;
    
    // Input/Time
    std::map<int, bool> keyStates;
    std::chrono::steady_clock::time_point lastFrameTime;
    float deltaTime = 0.016f;
    
    // Stats tracking
    std::chrono::steady_clock::time_point lastStatsUpdate;
    std::string cachedStats;
    bool statsNeedUpdate = true;

    enum class DebugColorMode {
        BASE,
        NUTRIENTS,
        MOISTURE
    };
    DebugColorMode currentColorMode = DebugColorMode::BASE;

public:
    worldboxSimUI() {
        // Position camera to look at the center of the world slightly from above
        cam.origin = v3(0, 50, 80);
        v3 target = v3(0, 0, 0);
        cam.direction = (target - cam.origin).normalized();
        cam.up = v3(0, 1, 0);
        cam.fov = 60;
        cam.rotationSpeed = M_1_PI;
        cam.movementSpeed = 50.0f;
        lastFrameTime = std::chrono::steady_clock::now();
    }

    ~worldboxSimUI() {
        if (textu != 0) {
            glDeleteTextures(1, &textu);
        }
        sim.grid.clear();
    }

    void renderUI(GLFWwindow* window) {
        // Compute delta time for consistent movement and sim stepping
        auto now = std::chrono::steady_clock::now();
        deltaTime = std::chrono::duration<float>(now - lastFrameTime).count();
        lastFrameTime = now;

        handleCameraControls(window);
        
        // Update simulation objects like the Star
        sim.updateStar(deltaTime);
        sim.updateWeatherAndPhysics(deltaTime);
        
        ImGui::Begin("WorldBox Simulation");
        if (ImGui::BeginTable("MainLayout", 2, ImGuiTableFlags_Resizable | ImGuiTableFlags_BordersOuter)) {
            ImGui::TableSetupColumn("Controls", ImGuiTableColumnFlags_WidthStretch, 0.3f);
            ImGui::TableSetupColumn("Preview", ImGuiTableColumnFlags_WidthStretch, 0.7f);
            ImGui::TableNextColumn();
            renderControlsPanel();
            ImGui::TableNextColumn();
            renderPreviewPanel();
            ImGui::EndTable();
        }
        ImGui::End();
    }

    void handleCameraControls(GLFWwindow* window) {
        glfwPollEvents();
        for (int i = GLFW_KEY_SPACE; i <= GLFW_KEY_LAST; i++) {
            keyStates[i] = (glfwGetKey(window, i) == GLFW_PRESS);
        }
        
        float currentSpeed = cam.movementSpeed * deltaTime;
        if (keyStates[GLFW_KEY_LEFT_SHIFT]) currentSpeed *= 3.0f; // Sprint

        if (keyStates[GLFW_KEY_W]) cam.moveForward(currentSpeed);
        if (keyStates[GLFW_KEY_S]) cam.moveBackward(currentSpeed);
        if (keyStates[GLFW_KEY_A]) cam.moveLeft(currentSpeed);
        if (keyStates[GLFW_KEY_D]) cam.moveRight(currentSpeed);
        if (keyStates[GLFW_KEY_E]) cam.moveUp(currentSpeed);
        if (keyStates[GLFW_KEY_Q]) cam.moveDown(currentSpeed);
        
        if (keyStates[GLFW_KEY_LEFT]) cam.rotateYaw(cam.rotationSpeed * deltaTime);
        if (keyStates[GLFW_KEY_RIGHT]) cam.rotateYaw(-cam.rotationSpeed * deltaTime);
        if (keyStates[GLFW_KEY_UP]) cam.rotatePitch(cam.rotationSpeed * deltaTime);
        if (keyStates[GLFW_KEY_DOWN]) cam.rotatePitch(-cam.rotationSpeed * deltaTime);
    }

    void renderControlsPanel() {
        ImGui::BeginChild("ControlsScroll", ImVec2(0, 0), true);
        
        if (ImGui::CollapsingHeader("World Generation", ImGuiTreeNodeFlags_DefaultOpen)) {
            ImGui::DragFloat("Width (X)", &sim.config.worldSizeX, 1.0f, 10.0f, 2000.0f);
            ImGui::DragFloat("Length (Z)", &sim.config.worldSizeZ, 1.0f, 10.0f, 2000.0f);
            ImGui::DragFloat("Depth (Y)", &sim.config.worldDepth, 1.0f, 1.0f, 500.0f);
            ImGui::DragFloat("Voxel Size", &sim.config.voxelSize, 0.1f, 0.1f, 10.0f);
            
            ImGui::ColorEdit3("Dirt Base Color", sim.config.baseDirtColor.data());
            
            ImGui::Separator();
            
            if (ImGui::Button("Generate Flat World", ImVec2(-1, 40))) {
                sim.generateFlatWorld();
                // applyDebugColorMode();
                statsNeedUpdate = true;
            }
            if (ImGui::Button("Clear World", ImVec2(-1, 30))) {
                sim.clearWorld();
                statsNeedUpdate = true;
            }
        }

        if (ImGui::CollapsingHeader("Weather & Physics", ImGuiTreeNodeFlags_DefaultOpen)) {
            ImGui::Checkbox("Enable Gravity (Terrain)", &sim.config.enableGravity);
            ImGui::DragFloat3("Gravity", sim.config.gravity.data());
            ImGui::DragFloat3("Wind", sim.config.wind.data());
            ImGui::DragFloat("Physics Step (sec)", &sim.config.physicsStep, 0.01f, 0.01f, 1.0f);
            
            ImGui::Separator();
            ImGui::Text("Clouds & Rain");
            ImGui::DragInt("Cloud Count", &sim.config.cloudCount, 1, 0, 100);
            ImGui::DragFloat("Cloud Height", &sim.config.cloudHeight, 5.0f, 10.0f, 1000.0f);
            ImGui::DragFloat("Rain Spawn Rate", &sim.config.rainSpawnRate, 0.1f, 0.0f, 50.0f);
            ImGui::ColorEdit3("Cloud Color", sim.config.cloudColor.data());
            ImGui::ColorEdit3("Rain Color", sim.config.rainColor.data());
            
            if (ImGui::Button("Generate Clouds", ImVec2(-1, 40))) {
                sim.generateClouds();
                applyDebugColorMode();
                statsNeedUpdate = true;
            }
            if (ImGui::Button("Clear Weather", ImVec2(-1, 30))) {
                for (auto& c : sim.clouds) sim.grid.remove(c.pos);
                for (auto& r : sim.rainDrops) sim.grid.remove(r.pos);
                sim.clouds.clear();
                sim.rainDrops.clear();
                statsNeedUpdate = true;
            }
        }

        if (ImGui::CollapsingHeader("Environment & Celestial", ImGuiTreeNodeFlags_DefaultOpen)) {
            ImGui::Text("Star Settings");
            ImGui::Checkbox("Enable Star Rotation", &sim.config.enableStarRotation);
            ImGui::DragFloat("Orbit Radius", &sim.config.starOrbitRadius, 10.0f, 100.0f, 5000.0f);
            ImGui::DragFloat("Star Speed", &sim.config.starSpeed, 0.01f, 0.0f, 5.0f);
            ImGui::DragFloat("Panel Size", &sim.config.starPanelSize, 5.0f, 10.0f, 1000.0f);
            ImGui::ColorEdit3("Star Color", sim.config.starColor.data());

            ImGui::Separator();
            ImGui::Text("Grass Settings");
            ImGui::SliderFloat("Grass Density", &sim.config.grassDensity, 0.0f, 1.0f);
            ImGui::ColorEdit3("Grass Color Base", sim.config.grassColorBase.data());
            
            if (ImGui::Button("Generate Grass", ImVec2(-1, 40))) {
                sim.generateGrass();
                applyDebugColorMode();
                statsNeedUpdate = true;
            }
        }

        if (ImGui::CollapsingHeader("Debug Views")) {
            ImGui::Text("Render Data Mode:");
            bool colorChanged = false;
            if (ImGui::RadioButton("Base Color", currentColorMode == DebugColorMode::BASE)) { 
                currentColorMode = DebugColorMode::BASE;
                colorChanged = true;
            }
            if (ImGui::RadioButton("Nutrients", currentColorMode == DebugColorMode::NUTRIENTS)) { 
                currentColorMode = DebugColorMode::NUTRIENTS;
                colorChanged = true;
            }
            if (ImGui::RadioButton("Moisture", currentColorMode == DebugColorMode::MOISTURE)) { 
                currentColorMode = DebugColorMode::MOISTURE;
                colorChanged = true;
            }

            if (colorChanged) {
                applyDebugColorMode();
            }
        }

        if (ImGui::CollapsingHeader("Camera & Render Settings", ImGuiTreeNodeFlags_DefaultOpen)) {
            ImGui::DragFloat3("Origin", cam.origin.data());
            ImGui::DragFloat3("Direction", cam.direction.data(), 0.0001f, -1.0f, 1.0f);
            ImGui::DragFloat("Movement Speed", &cam.movementSpeed, 0.1f, 1.0f, 500.0f);
            ImGui::InputFloat("Max Framerate", &framerate, 1, 10);
            
            ImGui::Checkbox("Use PBR/Raytracing (Slow)", &slowRender);
            if(slowRender) {
                ImGui::DragInt("Bounces", &reflectCount, 1, 0, 10);
                ImGui::Checkbox("Global Illumination", &globalIllumination);
            }
            
            ImGui::Checkbox("Use LODs", &useLod);

            if (ImGui::Button("Reset Camera")) {
                cam.origin = v3(0, sim.config.worldDepth * 2.0f, std::max(sim.config.worldSizeX, sim.config.worldSizeZ));
                cam.direction = (v3(0, 0, 0) - cam.origin).normalized();
            }
        }

        updateStatsCache();
        ImGui::TextUnformatted(cachedStats.c_str());

        ImGui::EndChild();
    }

    void renderPreviewPanel() {
        ImGui::BeginChild("PreviewChild", ImVec2(0, 0), true, ImGuiWindowFlags_NoScrollbar | ImGuiWindowFlags_NoScrollWithMouse);
        
        livePreview();
        if (textureInitialized) {
            float aspect = (float)currentPreviewFrame.getWidth() / (float)currentPreviewFrame.getHeight();
            float availWidth = ImGui::GetContentRegionAvail().x;
            ImGui::Image((void*)(intptr_t)textu, ImVec2(availWidth, availWidth / aspect));
        }

        ImGui::EndChild();
    }

    void applyDebugColorMode() {
        if (sim.grid.empty()) return;

        v3 boundsHalfExtent = v3(sim.config.worldSizeX, sim.config.worldDepth, sim.config.worldSizeZ);
        float searchRadius = boundsHalfExtent.norm() * 2.0f;

        auto allNodes = sim.grid.findInRadius(v3(0,0,0), searchRadius);

        for (auto& p : allNodes) {
            if (!p || !p->active) continue;

            v3 color = sim.config.baseDirtColor;
            
            switch (currentColorMode) {
                case DebugColorMode::NUTRIENTS: {
                    float t = std::clamp(p->data.nutrients, 0.0f, 1.0f);
                    color = v3(1.0f - t, t, 0.0f);
                    break;
                }
                case DebugColorMode::MOISTURE: {
                    float t = std::clamp(p->data.moisture, 0.0f, 1.0f);
                    color = v3(1.0f - t, 1.0f - t, 1.0f);
                    break;
                }
                case DebugColorMode::BASE:
                default:
                    if (p->data.type == 0) {
                        v3 darkDirt = sim.config.baseDirtColor * 0.4f;
                        color = sim.config.baseDirtColor * (1.0f - p->data.moisture) + darkDirt * p->data.moisture;
                    } else if (p->data.type == 1) {
                        color = sim.config.baseRockColor;
                    } else if (p->data.type == 2) {
                        v3 lushGrass = sim.config.grassColorBase * 1.5f;
                        color = sim.config.grassColorBase * (1.0f - p->data.moisture) + lushGrass * p->data.moisture;
                    } else if (p->data.type == 3) {
                        color = sim.config.starColor;
                    } else if (p->data.type == 4) {
                        color = sim.config.cloudColor;
                    } else if (p->data.type == 5) {
                        color = sim.config.rainColor;
                    }
                    break;
            }

            sim.grid.setColor(p->position, color);
        }
    }

    void livePreview() {
        std::lock_guard<std::mutex> lock(PreviewMutex);
        updatePreview = true;
        
        float invFrameRate = 1.0f / framerate;
        if (!useLod) {
            sim.grid.setLODFalloff(0.01);
            sim.grid.setLODMinDistance(1000.0f);
        } else {
            sim.grid.setLODMinDistance(lodDist);
            sim.grid.setLODFalloff(lodDropoff);
        }
        sim.grid.setMaxDistance(maxViewDistance);
        
        if (slowRender) {
            currentPreviewFrame = sim.grid.renderFrameTimed(cam, outHeight, outWidth, frame::colormap::RGB, invFrameRate, reflectCount, globalIllumination, useLod);
        } else {
            currentPreviewFrame = sim.grid.fastRenderFrame(cam, outHeight, outWidth, frame::colormap::RGB);
        }
        
        if (textu == 0) {
            glGenTextures(1, &textu);
        }
        glBindTexture(GL_TEXTURE_2D, textu);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
        
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, currentPreviewFrame.getWidth(), currentPreviewFrame.getHeight(), 
                     0, GL_RGB, GL_UNSIGNED_BYTE, currentPreviewFrame.getData().data());
        
        updatePreview = false;
        textureInitialized = true;
    }

    void updateStatsCache() {
        auto now = std::chrono::steady_clock::now();
        if (statsNeedUpdate || std::chrono::duration_cast<std::chrono::seconds>(now - lastStatsUpdate).count() >= 2) {
            std::stringstream gridstats;
            sim.grid.printStats(gridstats);
            cachedStats = gridstats.str();
            lastStatsUpdate = now;
            statsNeedUpdate = false;
        }
    }
};

#endif