#include "render.h"

#include <algorithm>
#include <cmath>

#include "raymath.h"

namespace hpp {

// One shader serves every geometry source. uMode: 0 = flat shaded, 1 = ambient + sun + point lights
// (cheap per-pixel Lambert), 2 = baked per-vertex lighting produced by hmap::bakeLighting().
static const char* kVS = R"GLSL(#version 330
in vec3 vertexPosition;
in vec2 vertexTexCoord;
in vec3 vertexNormal;
in vec4 vertexColor;
uniform mat4 mvp;
uniform mat4 matModel;
uniform mat4 matNormal;
out vec3 fragPos;
out vec2 fragUV;
out vec3 fragN;
out vec4 fragColor;
void main() {
    fragPos = vec3(matModel * vec4(vertexPosition, 1.0));
    fragUV = vertexTexCoord;
    fragN = normalize(vec3(matNormal * vec4(vertexNormal, 0.0)));
    fragColor = vertexColor;
    gl_Position = mvp * vec4(vertexPosition, 1.0);
}
)GLSL";

static const char* kFS = R"GLSL(#version 330
in vec3 fragPos;
in vec2 fragUV;
in vec3 fragN;
in vec4 fragColor;
uniform sampler2D texture0;
uniform vec4 colDiffuse;
uniform int uMode;
uniform vec3 uAmbient;
uniform vec3 uSunDir;
uniform vec3 uSunColor;
uniform int uLightCount;
uniform vec4 uLightPosR[16];
uniform vec4 uLightColI[16];
out vec4 finalColor;
void main() {
    vec4 tex = texture(texture0, fragUV) * colDiffuse;
    vec3 n = normalize(fragN);
    vec3 light;
    if (uMode == 1) {
        light = uAmbient + uSunColor * max(dot(n, -uSunDir), 0.0);
        for (int i = 0; i < uLightCount; i++) {
            vec3 d = uLightPosR[i].xyz - fragPos;
            float dist = length(d);
            float r = uLightPosR[i].w;
            float att = clamp(1.0 - (dist * dist) / (r * r), 0.0, 1.0);
            att *= att;
            light += uLightColI[i].rgb * uLightColI[i].a * max(dot(n, d / max(dist, 0.001)), 0.0) * att;
        }
    } else if (uMode == 2) {
        light = fragColor.rgb;
    } else {
        light = vec3(0.50 + 0.50 * max(dot(n, normalize(vec3(0.45, 0.8, 0.4))), 0.0));
    }
    finalColor = vec4(tex.rgb * light, tex.a);
}
)GLSL";

void Renderer::init() {
    shader_ = LoadShaderFromMemory(kVS, kFS);
    locMode_ = GetShaderLocation(shader_, "uMode");
    locAmbient_ = GetShaderLocation(shader_, "uAmbient");
    locSunDir_ = GetShaderLocation(shader_, "uSunDir");
    locSunColor_ = GetShaderLocation(shader_, "uSunColor");
    locCount_ = GetShaderLocation(shader_, "uLightCount");
    locPos_ = GetShaderLocation(shader_, "uLightPosR");
    locCol_ = GetShaderLocation(shader_, "uLightColI");
    mat_ = LoadMaterialDefault();
    mat_.shader = shader_;
}

void Renderer::freeGpu(std::vector<Gpu>& v) {
    for (Gpu& g : v) UnloadMesh(g.mesh);
    v.clear();
}

void Renderer::clearCaches() {
    for (auto& kv : textures_) UnloadTexture(kv.second);
    textures_.clear();
    for (auto& kv : models_) {
        if (!kv.second.ok) continue;
        Shader def = {rlGetShaderIdDefault(), rlGetShaderLocsDefault()};
        for (int i = 0; i < kv.second.model.materialCount; ++i) kv.second.model.materials[i].shader = def;  // keep our shader alive
        UnloadModel(kv.second.model);
    }
    models_.clear();
}

void Renderer::shutdown() {
    freeGpu(edit_);
    freeGpu(compiled_);
    clearCaches();
    MemFree(mat_.maps);
    mat_.maps = nullptr;
    UnloadShader(shader_);
}

void Renderer::upload(const std::vector<hmap::Batch>& in, std::vector<Gpu>& out, size_t& tris) {
    freeGpu(out);
    tris = 0;
    for (const hmap::Batch& b : in) {
        size_t n = b.indices.size();
        n -= n % 3;
        if (n == 0) continue;
        Mesh m = {};
        m.vertexCount = (int)n;  // expanded to a plain triangle list: no 16-bit index limit
        m.triangleCount = (int)(n / 3);
        m.vertices = (float*)MemAlloc((unsigned)(n * 3 * sizeof(float)));
        m.texcoords = (float*)MemAlloc((unsigned)(n * 2 * sizeof(float)));
        m.normals = (float*)MemAlloc((unsigned)(n * 3 * sizeof(float)));
        m.colors = (unsigned char*)MemAlloc((unsigned)(n * 4));
        for (size_t i = 0; i < n; ++i) {
            const hmap::Vertex& v = b.vertices[b.indices[i]];
            m.vertices[i * 3 + 0] = v.pos.x; m.vertices[i * 3 + 1] = v.pos.y; m.vertices[i * 3 + 2] = v.pos.z;
            m.normals[i * 3 + 0] = v.normal.x; m.normals[i * 3 + 1] = v.normal.y; m.normals[i * 3 + 2] = v.normal.z;
            m.texcoords[i * 2 + 0] = v.uv.x; m.texcoords[i * 2 + 1] = v.uv.y;
            m.colors[i * 4 + 0] = (unsigned char)(v.color & 0xff);
            m.colors[i * 4 + 1] = (unsigned char)((v.color >> 8) & 0xff);
            m.colors[i * 4 + 2] = (unsigned char)((v.color >> 16) & 0xff);
            m.colors[i * 4 + 3] = (unsigned char)((v.color >> 24) & 0xff);
        }
        UploadMesh(&m, false);
        out.push_back({m, b.material});
        tris += n / 3;
    }
}

void Renderer::setEditBatches(const std::vector<hmap::Batch>& b) { upload(b, edit_, editTriangles); }
void Renderer::setCompiledBatches(const std::vector<hmap::Batch>& b) { upload(b, compiled_, compiledTriangles); }

void Renderer::setLights(const std::vector<hmap::Light>& lights, const hmap::Vec3& ambient) {
    float amb[3] = {ambient.x, ambient.y, ambient.z};
    float sunDir[3] = {0, -1, 0}, sunCol[3] = {0, 0, 0};
    float pos[16 * 4] = {}, col[16 * 4] = {};
    int n = 0;
    bool haveSun = false;
    for (const hmap::Light& L : lights) {
        if (L.type == hmap::LightType::Sun) {
            if (haveSun) continue;
            hmap::Vec3 d = hmap::normalize(L.dir);
            sunDir[0] = d.x; sunDir[1] = d.y; sunDir[2] = d.z;
            sunCol[0] = L.color.x * L.intensity; sunCol[1] = L.color.y * L.intensity; sunCol[2] = L.color.z * L.intensity;
            haveSun = true;
            continue;
        }
        if (n >= 16) continue;  // spot lights are previewed as point lights
        pos[n * 4 + 0] = L.pos.x; pos[n * 4 + 1] = L.pos.y; pos[n * 4 + 2] = L.pos.z; pos[n * 4 + 3] = std::max(L.radius, 1.0f);
        col[n * 4 + 0] = L.color.x; col[n * 4 + 1] = L.color.y; col[n * 4 + 2] = L.color.z; col[n * 4 + 3] = L.intensity;
        ++n;
    }
    SetShaderValue(shader_, locAmbient_, amb, SHADER_UNIFORM_VEC3);
    SetShaderValue(shader_, locSunDir_, sunDir, SHADER_UNIFORM_VEC3);
    SetShaderValue(shader_, locSunColor_, sunCol, SHADER_UNIFORM_VEC3);
    SetShaderValue(shader_, locCount_, &n, SHADER_UNIFORM_INT);
    if (n > 0) {
        SetShaderValueV(shader_, locPos_, pos, SHADER_UNIFORM_VEC4, n);
        SetShaderValueV(shader_, locCol_, col, SHADER_UNIFORM_VEC4, n);
    }
}

void Renderer::setMode(int m) { SetShaderValue(shader_, locMode_, &m, SHADER_UNIFORM_INT); }

Texture2D Renderer::texture(const hmap::Material& m, const std::string& root) {
    std::string key = m.name + "|" + m.texture + "|" + root;
    auto it = textures_.find(key);
    if (it != textures_.end()) return it->second;
    Texture2D t = {};
    if (!m.texture.empty()) {
        std::string path = root + "/" + m.texture;
        if (FileExists(path.c_str())) t = LoadTexture(path.c_str());
    }
    if (t.id == 0) {  // no file: procedural checker, tinted by Material::color at draw time
        Image img = GenImageChecked(64, 64, 8, 8, Color{236, 236, 236, 255}, Color{176, 176, 182, 255});
        t = LoadTextureFromImage(img);
        UnloadImage(img);
    }
    GenTextureMipmaps(&t);
    SetTextureFilter(t, TEXTURE_FILTER_TRILINEAR);
    SetTextureWrap(t, TEXTURE_WRAP_REPEAT);
    textures_[key] = t;
    return t;
}

static bool isTranslucent(const std::string& name) {
    return name.find("trigger") != std::string::npos || name.find("clip") != std::string::npos ||
           name.find("water") != std::string::npos || name.find("glass") != std::string::npos;
}

void Renderer::drawBatches(bool useCompiled, const std::vector<hmap::Material>& mats, const std::string& root, int mode) {
    setMode(mode);
    static const hmap::Material fallback;
    std::vector<Gpu>& list = useCompiled ? compiled_ : edit_;
    for (int pass = 0; pass < 2; ++pass) {  // opaque first, then translucent volumes
        for (Gpu& g : list) {
            const hmap::Material& hm = (g.material >= 0 && (size_t)g.material < mats.size()) ? mats[(size_t)g.material] : fallback;
            bool tr = isTranslucent(hm.name);
            if ((pass == 1) != tr) continue;
            auto q = [](float v) { return (unsigned char)(std::max(0.0f, std::min(1.0f, v)) * 255.0f); };
            mat_.maps[MATERIAL_MAP_DIFFUSE].texture = texture(hm, root);
            mat_.maps[MATERIAL_MAP_DIFFUSE].color = Color{q(hm.color.x), q(hm.color.y), q(hm.color.z), (unsigned char)(tr ? 130 : 255)};
            DrawMesh(g.mesh, mat_, MatrixIdentity());
        }
    }
}

ModelEntry& Renderer::model(const std::string& rel, const std::string& root) {
    auto it = models_.find(rel);
    if (it != models_.end()) return it->second;
    ModelEntry me;
    bool isAbs = rel.size() > 1 && (rel[0] == '/' || rel[1] == ':');
    std::string path = isAbs ? rel : root + "/" + rel;
    if (!rel.empty() && FileExists(path.c_str())) {
        me.model = LoadModel(path.c_str());
        if (me.model.meshCount > 0) {
            me.ok = true;
            for (int i = 0; i < me.model.materialCount; ++i) {
                me.model.materials[i].shader = shader_;
                me.baseColors.push_back(me.model.materials[i].maps[MATERIAL_MAP_DIFFUSE].color);
            }
            me.box = GetModelBoundingBox(me.model);
        }
    }
    return models_.emplace(rel, std::move(me)).first->second;
}

void Renderer::drawModel(ModelEntry& me, const hmap::Mat34& xf, Color tint, int mode) {
    if (!me.ok) return;
    setMode(mode);
    Matrix t = toRlMatrix(xf);
    for (int i = 0; i < me.model.meshCount; ++i) {
        int mi = me.model.meshMaterial[i];
        Material& m = me.model.materials[mi];
        Color base = (size_t)mi < me.baseColors.size() ? me.baseColors[(size_t)mi] : WHITE;
        m.maps[MATERIAL_MAP_DIFFUSE].color = Color{(unsigned char)(base.r * tint.r / 255), (unsigned char)(base.g * tint.g / 255),
                                                   (unsigned char)(base.b * tint.b / 255), (unsigned char)(base.a * tint.a / 255)};
        DrawMesh(me.model.meshes[i], m, MatrixMultiply(me.model.transform, t));
    }
}

}  // namespace hpp