#include "shader.h"
#include "json.h"
#include "gl.h"

#include "objects/object.h"

static const char *version = "#version 330\n";
static const char *compat = \
  "#define max(x, y) max(y, x)\n"
  "#define frac fract\n"
  "#define log10 log\n"
  "#define CAST2(x) vec2(x)\n"
  "#define CAST3(x) vec3(x)\n"
  "#define CAST4(x) vec4(x)\n"
  "#define CAST3X3(x) mat3(x)\n"
  "#define saturate(x) clamp(x, 0.0, 1.0)\n"
  "#define lerp mix\n"
  "#define texSample2D texture2D\n"
  "#define texSample2DLod texture2DLod\n"
  "#define texture2DLod texture2D\n"
  "#define atan2 atan\n"
  "#define fmod(x, y) ((x)-(y)*trunc((x)/(y)))\n"
  "#define ddx dFdx\n"
  "#define ddy(x) dFdy(-(x))\n"
  "#define GLSL 1\n";

static const char *mul = "#define mul(x, y) ((y) * (x))\n";

namespace Mauri
{

static std::unordered_map<std::string, Shader *> shader_map;
static std::unordered_map<u64, RenderShader *> render_shader_map;

auto get_shader(const std::string &name) -> Shader *
{
    auto rshader = shader_map.find(name);
    if (rshader != shader_map.end()) return new Shader(rshader->second);
    Shader *shader = new Shader(name);
    shader_map[name] = shader;
    return shader;
}

extern auto unload_shaders() -> void
{
    for (auto &entry : render_shader_map)
    {
        delete entry.second;
    }
}

} // namespace Mauri

using namespace Mauri;

Shader::Shader(const std::string &name)
    : name(name)
{
    Asset *vs_asset = asset_manager()->get(name, VERTEX_SHADER);
    Asset *fs_asset = asset_manager()->get(name, FRAGMENT_SHADER);
    this->build_shader_source(vs_asset->as_string(), this->vs_source);
    this->build_shader_source(fs_asset->as_string(), this->fs_source);
}

Shader::Shader(Shader *origin)
{
    for (Uniform *uniform : origin->uniforms)
    {
        Uniform *copy = new Uniform();
        copy->name = uniform->name;
        copy->uname = uniform->uname;
        copy->type = uniform->type;
        copy->default_value = uniform->default_value;
        this->uniforms.push_back(copy);
    }
    this->origin_ = origin;
}

static auto simple_audio_bars_hack(std::string_view &line) -> const std::string_view
{
    if (line == "	uint barFreq1 = frequency % RESOLUTION;\r\n")
    {
        return "	uint barFreq1 = uint(mod(frequency, RESOLUTION));\r\n";
    }
    else if (line == "	uint barFreq2 = (barFreq1 + 1) % RESOLUTION;\r\n")
    {
        return "	uint barFreq2 = uint(mod((barFreq1 + 1u), RESOLUTION));\r\n";
    }
    else if (line == "	int bar = step(shapeCoord.y, 0.5 * lerp(g_BarBounds.x, g_BarBounds.y, lerp(barVolume1L, barVolume2L, smoothstep(0, 1, fract(frequency)))));\r\n")
    {
        return "	float bar = step(shapeCoord.y, 0.5 * lerp(g_BarBounds.x, g_BarBounds.y, lerp(barVolume1L, barVolume2L, smoothstep(0, 1, fract(frequency)))));\r\n";
    }
    else if (line == "	int bar = step(1 - shapeCoord.y, lerp(g_BarBounds.x, g_BarBounds.y, lerp(barVolume1, barVolume2, smoothstep(0, 1, fract(frequency)))));\r\n")
    {
        return "	float bar = step(1 - shapeCoord.y, lerp(g_BarBounds.x, g_BarBounds.y, lerp(barVolume1, barVolume2, smoothstep(0, 1, fract(frequency)))));\r\n";
    }
    else if (line == "	int bar = step(1 - shapeCoord.y, barHeight);\r\n")
    {
        return "	float bar = step(1 - shapeCoord.y, barHeight);\r\n";
    }
    else if (line == "	int barLeft = step(shapeCoord.y, barHeightLeft);\r\n")
    {
        return "	float barLeft = step(shapeCoord.y, barHeightLeft);\r\n";
    }
    else if (line == "	int barRight = step(1 - shapeCoord.y, barHeightRight);\r\n")
    {
        return "	float barRight = step(1 - shapeCoord.y, barHeightRight);\r\n";
    }
    else if (line == "	shapeCoord.x += endAngle - startAngle < 0.0;\r\n")
    {
        return "	shapeCoord.x += endAngle - int(startAngle < 0.0);\r\n";
    }
    else if (line == "	bar *= shapeCoord.x > 0.0 && shapeCoord.x * sign(endAngle - startAngle) < 1.0;\r\n")
    {
        return "	bar *= int(shapeCoord.x > 0.0 && shapeCoord.x * sign(endAngle - startAngle) < 1.0);\r\n";
    }
    return line;
}

static auto bokeh_blur_hack(std::string_view &line) -> const std::string_view
{
    if (line == "	depth *= (depth < blurPreciseLimit) * 6.0;\r\n")
    {
        return "	depth *= int(depth < blurPreciseLimit) * 6.0;\r\n";
    }
    return line;
}

static auto hlsl_hack(std::string_view &line) -> const std::string_view
{
    // https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-per-component-math#the-vector-type
    // HLSL would take the first components in these situations (I think).
    if (line == "	float pointer = g_PointerPosition.xy * u_pointerSpeed;\r\n")
    {
        return "	float pointer = g_PointerPosition.x * u_pointerSpeed;\r\n";
    }
    if (line == "vec3 albedo = texSample2D(g_Texture0, v_TexCoord.xy);\r\n")
    {
        return "vec3 albedo = texSample2D(g_Texture0, v_TexCoord.xy).xyz;\r\n";
    }
    if (line == "float scale = pow(length(abs(v_TexCoord - CAST2(u_offset)) * 1.0), 3.0) * u_scale;\r\n")
    {
        return "float scale = pow(length(abs(v_TexCoord.xy - CAST2(u_offset)) * 1.0), 3.0) * u_scale;\r\n";
    }
    return line;
}

static auto apply_hacks(std::string_view line) -> const std::string_view
{
    line = simple_audio_bars_hack(line);
    line = bokeh_blur_hack(line);
    line = hlsl_hack(line);
    return line;
}

auto Shader::build_shader_source(const std::string_view &isource, std::stringstream &out) -> void
{
    Combo current_preprocessor = Combo("", 0);

    std::stringstream prefix;

    auto cursor = isource.begin();
    auto end = isource.end();

    while (cursor != end)
    {
        auto end_of_line = std::find(cursor, end, '\n');
        if (end_of_line != end) {
            end_of_line++;
        }
        std::string_view raw_line = std::string_view(cursor, end_of_line - cursor);
        cursor += raw_line.length();
        const std::string_view line = apply_hacks(raw_line);

        bool cut_line = false;
        bool prepend_line = false;

        if (line.compare(0, 8, "#include") == 0)
        {
            size_t open = line.find('"') + 1;
            std::string_view include_name = line.substr(open, line.rfind('"') - open);

            Asset *asset = asset_manager()->get(std::string(include_name), SHADER);
            if (asset->type != ASSET_VOID)
            {
                std::stringstream include_source;
                this->build_shader_source(asset->as_string(), include_source);
                out << include_source.rdbuf();
            }

            cut_line = true;
        }
        else if (line.compare(0, 3, "#if") == 0 && (line.compare(0, 4, "#ifd") != 0 && line.compare(0, 4, "#ifn") != 0))
        {
            std::string_view name_start = line.substr(line.find(' ') + 1);
            size_t name_end = name_start.find(' ');
            if (name_end != std::string::npos)
            {
                current_preprocessor.name = name_start.substr(0, name_end);
                //std::string_view value_start = name_start.substr(name_end + 4);
                //std::string_view value = value_start.substr(0);
            }
            else
            {
                current_preprocessor.name = name_start.substr(0);
            }
        }
        else if (line.compare(0, 5, "#else") == 0)
        {
            // Change value.
        }
        else if (line.compare(0, 6, "#endif") == 0)
        {
            current_preprocessor.name = "";
        }
        else if (line.compare(0, 7, "uniform") == 0)
        {
            Uniform *uniform = new Uniform();

            std::string_view type_string = line.substr(line.find(' ') + 1);
            uniform->type = uniform_type_from_string(type_string);

            size_t open;

            if (uniform->type == TYPE_INVALID)
            {
                open = type_string.find(' ');
                uniform->precision = type_string.substr(0, open);
                type_string = type_string.substr(open + 1);
                uniform->type = uniform_type_from_string(type_string);
            }

            open = type_string.find(' ') + 1;
            uniform->uname = type_string.substr(open, type_string.find(';') - open);

            if (uniform->type == TYPE_SAMPLER2D)
            {
                this->texture_count++;
                s32 index = std::stoi(uniform->uname.substr(9));
                uniform->default_value[0] = index;
                prepend_line = true;
            }

            size_t combo_start = line.find("// ");
            if (combo_start != std::string::npos)
            {
                std::string_view combo_string = line.substr(combo_start + 2);
                const json &root = json::parse(combo_string, nullptr, false);
                if (!root.is_discarded())
                {
                    if (root.contains("material"))
                    {
                        uniform->name = root["material"];
                    }
                    else if (root.contains("label"))
                    {
                        uniform->name = root["label"];
                    }

                    switch (uniform->type)
                    {
                    case TYPE_SAMPLER2D:
                    {
                        if (root.contains("combo"))
                        {
                            this->texture_combos[(s32)uniform->default_value[0]] = root["combo"];
                        }
                        if (root.contains("default"))
                        {
                            this->textures[(s32)uniform->default_value[0]] = root["default"];
                        }
                        break;
                    }
                    case TYPE_FLOAT:
                    case TYPE_VEC4:
                    case TYPE_VEC3:
                    case TYPE_VEC2:
                        Parser::parse_value<vec4>(root, "default", &uniform->default_value, vec4(0.f));
                        break;
                    case TYPE_MAT4:
                    case TYPE_INVALID:
                    default:
                        break;
                    }
                }
            }

            this->uniforms.push_back(uniform);
        }
        else if (line.compare(0, 7, "varying") == 0 && name == "effects/foliagesway") // HACK
        {
            std::string_view type_string = line.substr(line.find(' ') + 1);
            size_t open = type_string.find(' ') + 1;
            std::string_view variable_name = type_string.substr(open, type_string.find(';') - open);
            bool duplicate = false;
            for (const Variable &variable : this->variables)
            {
                if (variable.name == variable_name)
                {
                    duplicate = true;
                    break;
                }
            }
            if (!duplicate)
            {
                this->variables.emplace_back(Variable{std::string(variable_name), uniform_type_from_string(type_string)});
            }
            cut_line = true;
        }
        else if (line.compare(0, 10, "// [COMBO]") == 0)
        {
            size_t open = line.find('{');
            std::string_view combo = line.substr(open, line.length() - open);
            const json &root = json::parse(combo, nullptr, false);
            if (!root.is_discarded())
            {
                if (root.contains("default"))
                {
                    this->combos.emplace_back(Combo(root["combo"], root["default"]));
                }
                else
                {
                    this->combos.emplace_back(Combo(root["combo"], 0));
                }
            }
            cut_line = true;
        }

        if (!cut_line)
        {
            if (prepend_line) prefix << line;
            else out << line;
        }
    }

    prefix << out.rdbuf();
    out = std::move(prefix);
}

auto Combo::to_string() const -> const std::string
{
    return "#define " + this->name + " " + std::to_string(this->value) + "\n";
}

auto Variable::to_string() const -> const std::string
{
    return "varying " + uniform_type_to_string(this->type) + " " + this->name + ";\n";
}

auto Shader::compile(const std::vector<Combo> &combo_overrides) -> RenderShader *
{
    std::vector<Combo> adjusted_combos = this->origin()->combos;

    for (const Combo &combo : combo_overrides)
    {
        bool add = true;
        for (Combo &shader_combo : adjusted_combos)
        {
            if (combo.name == shader_combo.name)
            {
                shader_combo.value = combo.value;
                add = false;
            }
        }
        if (add) adjusted_combos.push_back(combo);
    }

    std::stringstream combos_str;

    for (const Combo &combo : adjusted_combos)
    {
        combos_str << combo.to_string();
    }

    u64 hash = std::hash<std::string>{}(this->origin()->name + combos_str.str());

    auto rshader = render_shader_map.find(hash);
    if (rshader != render_shader_map.end()) return rshader->second;

    std::stringstream preamble;

    preamble << version;
    preamble << compat;

    preamble << combos_str.str();

    for (const Variable &variable : this->origin()->variables)
    {
        preamble << variable.to_string();
    }

    std::string vs_source = preamble.str() + mul + this->origin()->vs_source.str();
    std::string fs_source = preamble.str() + this->origin()->fs_source.str();

    return render_shader_map[hash] = new RenderShader(vs_source, fs_source);
}

auto Shader::bind(RenderShader *shader, Engine *engine, Object *object, mat4x4 *model,
    const std::array<Texture *, MAX_TEXTURES> *textures) -> void
{
    shader->bind();

    for (const Uniform *uniform : this->uniforms)
    {
        if (uniform->type != TYPE_MAT4 && uniform->type != TYPE_MAT3 &&
            uniform->uname.find('[') == std::string::npos)
        {
            shader->set_uniform(uniform->uname, uniform->type, glm::value_ptr(*uniform->value));
        }
    }

    if (textures)
    {
        for (u32 i = 0; i < this->origin()->texture_count; i++)
        {
            Texture *texture = (*textures)[i];
            if (!texture) continue;

            texture->update(engine->time);
            texture->bind(i);

            std::string index = std::to_string(i);
            std::string texture_resolution  = "g_Texture" + index + "Resolution";
            std::string texture_translation = "g_Texture" + index + "Translation";
            std::string texture_rotation    = "g_Texture" + index + "Rotation";
            shader->set_uniform(texture_resolution, TYPE_VEC4, glm::value_ptr(*texture->texture_resolution));
            shader->set_uniform(texture_translation, TYPE_VEC2, glm::value_ptr(texture->texture_translation));
            shader->set_uniform(texture_rotation, TYPE_VEC4, glm::value_ptr(texture->texture_rotation));
        }
    }

    f32 adjusted_time = engine->time * engine->rate;
    shader->set_uniform("g_Time", TYPE_FLOAT, &adjusted_time);

    shader->set_uniform("g_TexelSize", TYPE_VEC2, glm::value_ptr(engine->view.texel_size));
    shader->set_uniform("g_TexelSizeHalf", TYPE_VEC2, glm::value_ptr(engine->view.texel_half_size));

    if (model)
    {
        mat4x4 inverse = glm::inverse(*model);
        shader->set_uniform("g_ModelViewProjectionMatrix", TYPE_MAT4, glm::value_ptr(*model));
        shader->set_uniform("g_ModelViewProjectionMatrixInverse", TYPE_MAT4, glm::value_ptr(inverse));
        shader->set_uniform("g_EffectModelViewProjectionMatrix", TYPE_MAT4, glm::value_ptr(engine->default_mat));
        shader->set_uniform("g_EffectModelViewProjectionMatrixInverse", TYPE_MAT4, glm::value_ptr(engine->default_mat_inverse));
        shader->set_uniform("g_EffectTextureProjectionMatrix", TYPE_MAT4, glm::value_ptr(engine->default_mat));
        shader->set_uniform("g_EffectTextureProjectionMatrixInverse", TYPE_MAT4, glm::value_ptr(engine->default_mat_inverse));

        shader->set_uniform("g_ModelMatrix", TYPE_MAT4, glm::value_ptr(engine->default_mat));
    }

    if (object)
    {
        vec2 cursor_pos = engine->view.cursor_pos;
        if (object->passthrough)
        {
            cursor_pos[0] -= object->size[0] / 2.f;
            cursor_pos[1] -= object->size[1] / 2.f;
        }
        cursor_pos[0] /= engine->view.width;
        cursor_pos[1] /= engine->view.height;
        //cursor_pos[1] = -cursor_pos[1];
        shader->set_uniform("g_PointerPosition", TYPE_VEC2, glm::value_ptr(cursor_pos));
        shader->set_uniform("g_Alpha", TYPE_FLOAT, &object->alpha);
        shader->set_uniform("g_Color", TYPE_VEC3, glm::value_ptr(object->color));
        //shader->set_uniform("g_Color4", TYPE_VEC4, glm::value_ptr(object->color4));
    }

    static const std::array<f32, 64> empty_bars = {};

    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum64Left"), 64, empty_bars.data());
    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum64Right"), 64, empty_bars.data());

    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum32Left"), 32, empty_bars.data());
    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum32Right"), 32, empty_bars.data());

    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum16Left"), 16, empty_bars.data());
    glUniform1fv(shader->get_uniform_location("g_AudioSpectrum16Right"), 16, empty_bars.data());
}

Shader::~Shader()
{
    for (Uniform *uniform : this->uniforms)
    {
        delete uniform;
    }
}