#ifndef WORLDBOX_HPP
#define WORLDBOX_HPP

#include <iostream>
#include <vector>
#include <chrono>
#include <thread>
#include <mutex>
#include <cmath>
#include <random>
#include <algorithm>

#include "../grid/grid3eigen.hpp"
#include "../timing_decorator.cpp"

using v3 = Eigen::Vector3f;

struct WorldVoxel {
    float nutrients = 1.0f;
    float moisture = 0.5f;
    int type = 0; // 0=Dirt, 1=Rock, 2=Grass, 3=Star, 4=Cloud, 5=Rain

    WorldVoxel() = default;
    
    WorldVoxel(float nut, float mois, int t) : nutrients(nut), moisture(mois), type(t) {}
};

struct WorldBoxConfig {
    v3 center = v3(0, 0, 0);
    float worldSizeX = 1000.0f;
    float worldSizeZ = 1000.0f;
    float worldDepth = 20.0f;
    
    float voxelSize = 2.0f;
    v3 baseDirtColor = v3(0.36f, 0.25f, 0.14f);
    v3 baseRockColor = v3(0.45f, 0.45f, 0.45f);
    
    float gridSizeCubeMin = 1024.0f;

    // Grass Config
    float grassDensity = 0.05f;
    v3 grassColorBase = v3(0.2f, 0.6f, 0.15f);

    // Star Config
    bool enableStarRotation = false; // Off by default
    float starOrbitRadius = 800.0f;
    float starPanelSize = 100.0f;
    float starVoxelSize = 10.0f;
    v3 starColor = v3(1.0f, 0.95f, 0.8f);
    float starSpeed = 0.2f; // Radians per second
    float starAngle = 0.0f; 

    // Weather Config
    int cloudCount = 15;
    float cloudHeight = 150.0f;
    v3 cloudColor = v3(0.9f, 0.9f, 0.95f);
    float cloudBaseSize = 6.0f;

    v3 rainColor = v3(0.2f, 0.4f, 0.9f);
    float rainDropSize = 0.5f;
    float rainSpawnRate = 1.0f;

    // Physics Config
    bool enableGravity = true;
    v3 gravity = v3(0.0f, -60.0f, 0.0f);
    v3 wind = v3(20.0f, 0.0f, 10.0f);
    float physicsStep = 0.1f;
};

struct CloudVoxel {
    v3 pos;
    float size;
};

struct RainDrop {
    v3 pos;
    v3 vel;
};

class worldboxsim {
public:
    WorldBoxConfig config;
    Octree<WorldVoxel> grid;
    std::mt19937 rng;
    std::vector<v3> starVoxelPositions;
    
    std::vector<CloudVoxel> clouds;
    std::vector<RainDrop> rainDrops;
    float physicsTimer = 0.0f;

    worldboxsim() : rng(42) {
        config = WorldBoxConfig();
        grid = Octree<WorldVoxel>(v3(-config.gridSizeCubeMin, -config.gridSizeCubeMin, -config.gridSizeCubeMin),v3(config.gridSizeCubeMin,  config.gridSizeCubeMin,  config.gridSizeCubeMin), 16, 32);
        grid.setBackgroundColor(v3(0.53f, 0.81f, 0.92f));
    }

    void updateStar(float dt) {
        if (!config.enableStarRotation) {
            if (!starVoxelPositions.empty()) {
                for(const auto& pos : starVoxelPositions) {
                    grid.remove(pos); 
                }
                starVoxelPositions.clear();
            }
            return;
        }

        // Calculate rotation
        config.starAngle += dt * config.starSpeed;
        if (config.starAngle > 2 * M_PI) config.starAngle -= 2 * M_PI;

        // Calculate new center of star (orbiting on the X/Y plane)
        v3 starCenter(cos(config.starAngle) * config.starOrbitRadius, sin(config.starAngle) * config.starOrbitRadius, 0.0f);

        // Create a flat panel facing the origin
        v3 n = -starCenter.normalized();
        v3 worldUp(0, 1, 0);
        if (std::abs(n.dot(worldUp)) > 0.99f) worldUp = v3(0, 0, 1);
        v3 right = worldUp.cross(n).normalized();
        v3 up = n.cross(right).normalized();

        int halfGrid = std::max(1, static_cast<int>((config.starPanelSize / config.starVoxelSize) / 2.0f));
        WorldVoxel starVoxel(0.0f, 0.0f, 3); // Type 3 = Star
        
        // Calculate the new ideal positions for this frame
        std::vector<v3> newPositions;
        newPositions.reserve((2 * halfGrid + 1) * (2 * halfGrid + 1));

        for (int i = -halfGrid; i <= halfGrid; ++i) {
            for (int j = -halfGrid; j <= halfGrid; ++j) {
                newPositions.push_back(starCenter + (right * (i * config.starVoxelSize)) + (up * (j * config.starVoxelSize)));
            }
        }

        // Apply grid changes
        if (starVoxelPositions.empty()) {
            // Creation: Spawn voxels into the grid for the first time
            for (const auto& pos : newPositions) {
                // Injecting a high emittance factor (15.0f) to make it a bright emissive light source
                grid.set(starVoxel, pos, true, config.starColor, config.starVoxelSize, true, 1, 1, 15.0f);
            }
            starVoxelPositions = newPositions;
        } else if (starVoxelPositions.size() == newPositions.size()) {
            // Moving: Using grid.move() to smoothly transfer nodes in the Octree
            for (size_t i = 0; i < starVoxelPositions.size(); ++i) {
                grid.move(starVoxelPositions[i], newPositions[i]);
            }
            starVoxelPositions = newPositions;
        } 
    }

    void generateClouds() {
        std::uniform_real_distribution<float> randX(-config.worldSizeX/2, config.worldSizeX/2);
        std::uniform_real_distribution<float> randZ(-config.worldSizeZ/2, config.worldSizeZ/2);
        std::uniform_real_distribution<float> randY(config.cloudHeight - 10.0f, config.cloudHeight + 10.0f);

        for (int i=0; i<config.cloudCount; ++i) {
            v3 center(randX(rng), randY(rng), randZ(rng));
            int numVoxels = 10 + (rng() % 20);
            
            for (int j=0; j<numVoxels; ++j) {
                v3 offset(
                    (rng() % 40) - 20,
                    (rng() % 10) - 5,
                    (rng() % 40) - 20
                );
                v3 pos = center + offset;
                float size = config.cloudBaseSize + (rng() % 6);
                
                WorldVoxel vox(0.0f, 1.0f, 4); // Type 4 = Cloud
                // Adding to grid with transmission=0.4f (makes it partially transparent for RTX)
                grid.set(vox, pos, true, config.cloudColor, size, true, 4, 0, 0.0f, 1.0f, 0.0f, 0.4f);
                clouds.push_back({pos, size});
            }
        }
    }

    void updateWeatherAndPhysics(float dt) {
        float halfX = config.worldSizeX / 2.0f;
        float halfZ = config.worldSizeZ / 2.0f;

        // 1. Clouds Update
        std::vector<CloudVoxel> nextClouds;
        for (auto& c : clouds) {
            v3 nextPos = c.pos + config.wind * dt;
            
            // Screen wrap logic for wind drift
            if (nextPos.x() > halfX) nextPos.x() -= config.worldSizeX;
            if (nextPos.x() < -halfX) nextPos.x() += config.worldSizeX;
            if (nextPos.z() > halfZ) nextPos.z() -= config.worldSizeZ;
            if (nextPos.z() < -halfZ) nextPos.z() += config.worldSizeZ;
            
            if (grid.move(c.pos, nextPos)) {
                c.pos = nextPos;
            } else {
                WorldVoxel vox(0.0f, 1.0f, 4);
                grid.set(vox, nextPos, true, config.cloudColor, c.size, true, 4, 0, 0.0f, 1.0f, 0.0f, 0.4f);
                c.pos = nextPos;
            }
            nextClouds.push_back(c);
            
            // Spawn Rain
            std::uniform_real_distribution<float> dist(0, 1);
            if (dist(rng) < (config.rainSpawnRate * dt * 0.1f)) {
                RainDrop r = {c.pos - v3(0, c.size, 0), config.wind};
                rainDrops.push_back(r);
                WorldVoxel rv(0.0f, 1.0f, 5); // Type 5 = Rain
                grid.set(rv, r.pos, true, config.rainColor, config.rainDropSize, true, 5);
            }
        }
        clouds = nextClouds;
        
        // 2. Rain Update
        std::vector<RainDrop> nextRain;
        for (auto& r : rainDrops) {
            r.vel += config.gravity * dt;
            v3 nextPos = r.pos + r.vel * dt;
            
            v3 dir = (nextPos - r.pos);
            float distMag = dir.norm();
            
            if (distMag > 0) {
                dir.normalize();
                auto hit = grid.voxelTraverse(r.pos, dir, distMag, false);
                
                // If it hits solid terrain
                if (hit && hit->data.type != 4 && hit->data.type != 5) {
                    if (hit->data.type == 0) { // Hit Dirt
                        hit->data.moisture = std::min(1.0f, hit->data.moisture + 0.15f);
                        v3 darkDirt = config.baseDirtColor * 0.4f;
                        v3 wetColor = config.baseDirtColor * (1.0f - hit->data.moisture) + darkDirt * hit->data.moisture;
                        grid.setColor(hit->position, wetColor);
                    } else if (hit->data.type == 2) { // Hit Grass
                        hit->data.moisture = std::min(1.0f, hit->data.moisture + 0.15f);
                        v3 lushGrass = config.grassColorBase * 1.5f;
                        v3 wetColor = config.grassColorBase * (1.0f - hit->data.moisture) + lushGrass * hit->data.moisture;
                        grid.setColor(hit->position, wetColor);
                    }
                    
                    grid.remove(r.pos);
                    continue; 
                }
            }
            
            // Delete if falls out of bounds
            if (nextPos.y() < -config.worldDepth - 20.0f) {
                grid.remove(r.pos);
                continue;
            }
            
            if (grid.move(r.pos, nextPos)) {
                r.pos = nextPos;
                nextRain.push_back(r);
            } else {
                WorldVoxel rv(0.0f, 1.0f, 5);
                grid.set(rv, nextPos, true, config.rainColor, config.rainDropSize, true, 5);
                r.pos = nextPos;
                nextRain.push_back(r);
            }
        }
        rainDrops = nextRain;
        
        // 3. Apply Block Gravity
        if (config.enableGravity) {
            physicsTimer += dt;
            if (physicsTimer >= config.physicsStep) {
                applyTerrainGravity();
                physicsTimer = 0.0f;
            }
        }
    }

    void applyTerrainGravity() {
        std::vector<std::shared_ptr<Octree<WorldVoxel>::NodeData>> nodes;
        grid.collectNodesByObjectId( -1, nodes);
        
        std::vector<std::shared_ptr<Octree<WorldVoxel>::NodeData>> terrain;
        terrain.reserve(nodes.size());
        for (auto& n : nodes) {
            // Include Dirt, Rock, and Grass in gravity sweep
            if (n->data.type == 0 || n->data.type == 1 || n->data.type == 2) {
                terrain.push_back(n);
            }
        }
        
        // Process Bottom-Up
        std::sort(terrain.begin(), terrain.end(), [](const auto& a, const auto& b) {
            return a->position.y() < b->position.y();
        });
        
        for (auto& n : terrain) {
            v3 belowPos = n->position + v3(0, -config.voxelSize, 0);
            
            // Bounds check so voxels don't fall infinitely
            if (belowPos.y() < -config.worldDepth) continue;
            
            auto hit = grid.find(belowPos, config.voxelSize * 0.1f);
            if (!hit) {
                grid.move(n->position, belowPos);
            }
        }
    }

    void generateGrass() {
        TIME_FUNCTION;
        float halfX = config.worldSizeX / 2.0f;
        float halfZ = config.worldSizeZ / 2.0f;
        float surfaceY = 0.0f; 

        int stepsX = static_cast<int>(std::round(config.worldSizeX / config.voxelSize)) + 1;
        int stepsZ = static_cast<int>(std::round(config.worldSizeZ / config.voxelSize)) + 1;

        int grassCount = 0;

        #pragma omp parallel
        {
            std::random_device rd;
            std::mt19937 local_rng(rd() ^ std::hash<std::thread::id>()(std::this_thread::get_id()));
            std::uniform_real_distribution<float> probDist(0.0f, 1.0f);
            std::uniform_int_distribution<int> grassHeightDist(1, 8);

            #pragma omp for schedule(static) collapse(2)
            for (int i = 0; i < stepsX; ++i) {
                for (int j = 0; j < stepsZ; ++j) {
                    
                    float x = -halfX + i * config.voxelSize;
                    float z = -halfZ + j * config.voxelSize;

                    if (x > halfX || z > halfZ) continue;
                    
                    if (probDist(local_rng) < config.grassDensity) {
                        int gHeight = grassHeightDist(local_rng);
                        float gSize = config.voxelSize / 25.0f;
                        
                        std::uniform_real_distribution<float> offDist(-config.voxelSize/2.0f + gSize/2.0f, config.voxelSize/2.0f - gSize/2.0f);
                        float offsetX = offDist(local_rng);
                        float offsetZ = offDist(local_rng);
                        
                        WorldVoxel gVox(1.0f, 0.8f, 2); // Type 2 = Grass
                        float baseY = surfaceY + (config.voxelSize / 2.0f) + (gSize / 2.0f);
                        
                        #pragma omp critical
                        {
                            for (int g = 0; g < gHeight; ++g) {
                                v3 gPos(x + offsetX, baseY + g * gSize, z + offsetZ);
                                grid.set(gVox, gPos, true, config.grassColorBase, gSize, true, 1, 0);
                                grassCount++;
                            }
                        }
                    }
                }
            }
        }
        std::cout << "Grass generation complete. Placed " << grassCount << " grass voxels." << std::endl;
    }

    void generateFlatWorld() {
        TIME_FUNCTION;
        grid.clear();
        
        float halfX = config.worldSizeX / 2.0f;
        float halfZ = config.worldSizeZ / 2.0f;
        float surfaceY = 0.0f; 

        // 1. Calculate integer bounds to satisfy OpenMP
        int stepsX = static_cast<int>(std::round(config.worldSizeX / config.voxelSize)) + 1;
        int stepsZ = static_cast<int>(std::round(config.worldSizeZ / config.voxelSize)) + 1;
        int stepsY = static_cast<int>(std::round(config.worldDepth / config.voxelSize)) + 1;
        size_t maxSteps = stepsX * stepsZ * stepsY;

        int nodeCount = 0;

        #pragma omp parallel for schedule(static) collapse(3)
        for (int i = 0; i < stepsX; ++i) {
            for (int j = 0; j < stepsZ; ++j) {
                for (int k = 0; k < stepsY; ++k) {
                    
                    float x = -halfX + i * config.voxelSize;
                    float z = -halfZ + j * config.voxelSize;
                    float y = surfaceY - k * config.voxelSize;

                    if (x > halfX || z > halfZ || y < surfaceY - config.worldDepth) {
                        continue;
                    }
                    
                    WorldVoxel voxel;
                    v3 color;
                    
                    float depthRatio = std::abs(y - surfaceY) / config.worldDepth;
                    
                    if (depthRatio > 0.8f) {
                        voxel.type = 1;
                        voxel.nutrients = 0.1f;
                        color = config.baseRockColor;
                    } else {
                        voxel.type = 0;
                        voxel.nutrients = 1.0f - depthRatio;
                        color = config.baseDirtColor;
                    }

                    v3 pos(x, y, z);
                
                    #pragma omp critical
                    grid.set(voxel, pos, true, color, config.voxelSize, true, 1, 0);
                    // nodeCount++;
                }
            }
        }
        
        std::cout << "World generation complete. Placed " << nodeCount << " voxels." << std::endl;
    }

    void clearWorld() {
        grid.clear();
        clouds.clear();
        rainDrops.clear();
        starVoxelPositions.clear();
    }
};

#endif