#pragma once

#include <raylib.h>
#include <raymath.h>
#include <vector>
#include <fstream>
#include <cstdint>

class Heightmap {
public:
    struct Config {
        float unitsPerSample = 1.0f;  // How many world units (meters) per heightmap sample
        float heightScale = 1.0f;      // Scale factor for height values
    };

private:
    std::vector<float> data;
    int samplesPerSide{0};        // Number of samples along one edge
    float worldWidth{0.0f};        // Total world width in units
    float worldHeight{0.0f};       // Total world height in units
    Config config;

public:
    Heightmap() : config{} {}
    Heightmap(const Config& cfg) : config(cfg) {}

    // Load heightmap from binary file
    bool load(const std::string& filename) {
        std::ifstream file(filename, std::ios::binary);
        if (!file) return false;

        int32_t fileSize;
        file.read(reinterpret_cast<char*>(&fileSize), sizeof(fileSize));
        samplesPerSide = fileSize;

        data.resize(samplesPerSide * samplesPerSide);
        file.read(reinterpret_cast<char*>(data.data()), data.size() * sizeof(float));

        // Calculate world dimensions based on samples and units per sample
        worldWidth = worldHeight = (samplesPerSide - 1) * config.unitsPerSample;

        return true;
    }

    // Get interpolated height at world position
    float getHeightAtPosition(float worldX, float worldZ) const {
        // Convert world coordinates to sample coordinates
        float sampleX = (worldX + worldWidth * 0.5f) / config.unitsPerSample;
        float sampleZ = (worldZ + worldHeight * 0.5f) / config.unitsPerSample;

        // Clamp to valid range
        sampleX = std::max(0.0f, std::min(sampleX, float(samplesPerSide - 1)));
        sampleZ = std::max(0.0f, std::min(sampleZ, float(samplesPerSide - 1)));

        // Get integer sample indices
        int x0 = static_cast<int>(sampleX);
        int z0 = static_cast<int>(sampleZ);
        int x1 = std::min(x0 + 1, samplesPerSide - 1);
        int z1 = std::min(z0 + 1, samplesPerSide - 1);

        // Get fractional parts for interpolation
        float fx = sampleX - x0;
        float fz = sampleZ - z0;

        // Get heights at four corners
        float h00 = data[z0 * samplesPerSide + x0] * config.heightScale;
        float h10 = data[z0 * samplesPerSide + x1] * config.heightScale;
        float h01 = data[z1 * samplesPerSide + x0] * config.heightScale;
        float h11 = data[z1 * samplesPerSide + x1] * config.heightScale;

        // Bilinear interpolation
        float h0 = h00 * (1.0f - fx) + h10 * fx;
        float h1 = h01 * (1.0f - fx) + h11 * fx;
        return h0 * (1.0f - fz) + h1 * fz;
    }

    // Generate mesh from heightmap
    Mesh generateMesh() const {
        auto mesh = Mesh{};
        int vertexCount = samplesPerSide * samplesPerSide;
        int triangleCount = (samplesPerSide - 1) * (samplesPerSide - 1) * 2;

        mesh.vertexCount = vertexCount;
        mesh.triangleCount = triangleCount;
        mesh.vertices = (float*)MemAlloc(vertexCount * 3 * sizeof(float));
        mesh.texcoords = (float*)MemAlloc(vertexCount * 2 * sizeof(float));
        mesh.normals = (float*)MemAlloc(vertexCount * 3 * sizeof(float));
        mesh.indices = (unsigned short*)MemAlloc(triangleCount * 3 * sizeof(unsigned short));

        // Generate vertices
        for (int z = 0; z < samplesPerSide; z++) {
            for (int x = 0; x < samplesPerSide; x++) {
                int idx = z * samplesPerSide + x;

                // World position in units (meters)
                // Center the terrain around origin
                float worldX = (x * config.unitsPerSample) - worldWidth * 0.5f;
                float worldZ = (z * config.unitsPerSample) - worldHeight * 0.5f;
                float worldY = data[idx] * config.heightScale;

                mesh.vertices[idx * 3] = worldX;
                mesh.vertices[idx * 3 + 1] = worldY;
                mesh.vertices[idx * 3 + 2] = worldZ;

                // Texture coordinates: 1 unit = 1 texture tile
                // Map from world coordinates to texture coordinates
                float texU = worldX + worldWidth * 0.5f;
                float texV = worldZ + worldHeight * 0.5f;

                mesh.texcoords[idx * 2] = texU;
                mesh.texcoords[idx * 2 + 1] = texV;
            }
        }

        // Generate indices
        int triIdx = 0;
        for (int z = 0; z < samplesPerSide - 1; z++) {
            for (int x = 0; x < samplesPerSide - 1; x++) {
                int topLeft = z * samplesPerSide + x;
                int topRight = topLeft + 1;
                int bottomLeft = (z + 1) * samplesPerSide + x;
                int bottomRight = bottomLeft + 1;

                mesh.indices[triIdx++] = topLeft;
                mesh.indices[triIdx++] = bottomLeft;
                mesh.indices[triIdx++] = topRight;

                mesh.indices[triIdx++] = topRight;
                mesh.indices[triIdx++] = bottomLeft;
                mesh.indices[triIdx++] = bottomRight;
            }
        }

        // Calculate proper normals
        calculateNormals(mesh);

        UploadMesh(&mesh, false);
        return mesh;
    }

    // Getters
    float getWorldWidth() const { return worldWidth; }
    float getWorldHeight() const { return worldHeight; }
    Vector3 getWorldBounds() const {
        return {worldWidth, getMaxHeight() * config.heightScale, worldHeight};
    }
    Vector3 getWorldCenter() const { return {0, 0, 0}; }
    int getSamplesPerSide() const { return samplesPerSide; }

private:
    void calculateNormals(Mesh& mesh) const {
        // Initialize normals to zero
        for (int i = 0; i < mesh.vertexCount * 3; i++) {
            mesh.normals[i] = 0.0f;
        }

        // Calculate face normals and add to vertex normals
        for (int i = 0; i < mesh.triangleCount; i++) {
            unsigned short i1 = mesh.indices[i * 3];
            unsigned short i2 = mesh.indices[i * 3 + 1];
            unsigned short i3 = mesh.indices[i * 3 + 2];

            Vector3 v1 = {mesh.vertices[i1 * 3], mesh.vertices[i1 * 3 + 1], mesh.vertices[i1 * 3 + 2]};
            Vector3 v2 = {mesh.vertices[i2 * 3], mesh.vertices[i2 * 3 + 1], mesh.vertices[i2 * 3 + 2]};
            Vector3 v3 = {mesh.vertices[i3 * 3], mesh.vertices[i3 * 3 + 1], mesh.vertices[i3 * 3 + 2]};

            Vector3 edge1 = Vector3Subtract(v2, v1);
            Vector3 edge2 = Vector3Subtract(v3, v1);
            Vector3 normal = Vector3Normalize(Vector3CrossProduct(edge1, edge2));

            // Add face normal to each vertex
            mesh.normals[i1 * 3] += normal.x;
            mesh.normals[i1 * 3 + 1] += normal.y;
            mesh.normals[i1 * 3 + 2] += normal.z;

            mesh.normals[i2 * 3] += normal.x;
            mesh.normals[i2 * 3 + 1] += normal.y;
            mesh.normals[i2 * 3 + 2] += normal.z;

            mesh.normals[i3 * 3] += normal.x;
            mesh.normals[i3 * 3 + 1] += normal.y;
            mesh.normals[i3 * 3 + 2] += normal.z;
        }

        // Normalize vertex normals
        for (int i = 0; i < mesh.vertexCount; i++) {
            Vector3 normal = {
                mesh.normals[i * 3],
                mesh.normals[i * 3 + 1],
                mesh.normals[i * 3 + 2]
            };
            normal = Vector3Normalize(normal);
            mesh.normals[i * 3] = normal.x;
            mesh.normals[i * 3 + 1] = normal.y;
            mesh.normals[i * 3 + 2] = normal.z;
        }
    }

    float getMaxHeight() const {
        float maxHeight = 0.0f;
        for (float h : data) {
            maxHeight = std::max(maxHeight, h);
        }
        return maxHeight;
    }
};