#include "materialcompiler.hpp"

#include "file.hpp"
#include "log.hpp"
#include "engine.hpp"
#include "string_utils.hpp"
#include "shadercompiler.hpp"
#include "renderer.hpp"

ShaderSource get_shader(const std::string& filename, bool skinned, bool cubemap) {
    auto shader_file = prism::open_file(prism::base_domain / filename);
    if(!shader_file.has_value()) {
        prism::log("Failed to open shader file {}!", filename);
        return {};
    }

    ShaderStage stage;
    if(filename.find("vert") != std::string::npos) {
        stage = ShaderStage::Vertex;
    } else {
        stage = ShaderStage::Fragment;
    }

    CompileOptions options;
    if(skinned)
        options.add_definition("BONE");

    if(cubemap)
        options.add_definition("CUBEMAP");

    return *shader_compiler.compile(ShaderLanguage::GLSL, stage, ShaderSource(shader_file->read_as_string()), engine->get_gfx()->accepted_shader_language(), options);
}

GFXPipeline* MaterialCompiler::create_static_pipeline(GFXGraphicsPipelineCreateInfo createInfo, bool positions_only, bool cubemap) {
    // take vertex src
    const std::string vertex_path = fmt::format("{}.glsl", createInfo.shaders.vertex_src.as_path().string());

    if (positions_only)
        createInfo.label += " (shadow)";

    if (cubemap)
        createInfo.label += " (cubemap)";

    createInfo.shaders.vertex_src = get_shader(vertex_path, false, cubemap);

    if(positions_only) {
        createInfo.vertex_input.inputs = {
            {position_buffer_index, sizeof(prism::float3)}
        };

        createInfo.vertex_input.attributes = {
            {position_buffer_index, 0, 0, GFXVertexFormat::FLOAT3}
        };
    } else {
        createInfo.vertex_input.inputs = {
            {position_buffer_index, sizeof(prism::float3)},
            {normal_buffer_index, sizeof(prism::float3)},
            {texcoord_buffer_index, sizeof(prism::float2)},
            {tangent_buffer_index, sizeof(prism::float3)},
            {bitangent_buffer_index, sizeof(prism::float3)}
        };

        createInfo.vertex_input.attributes = {
            {position_buffer_index, 0, 0, GFXVertexFormat::FLOAT3},
            {normal_buffer_index, 1, 0, GFXVertexFormat::FLOAT3},
            {texcoord_buffer_index, 2, 0, GFXVertexFormat::FLOAT2},
            {tangent_buffer_index, 3, 0, GFXVertexFormat::FLOAT3},
            {bitangent_buffer_index, 4, 0, GFXVertexFormat::FLOAT3}
        };
    }

    return engine->get_gfx()->create_graphics_pipeline(createInfo);
}

GFXPipeline* MaterialCompiler::create_skinned_pipeline(GFXGraphicsPipelineCreateInfo createInfo, bool positions_only) {
    createInfo.label += " (skinned)";

    // take vertex src
    const std::string vertex_path = fmt::format("{}.glsl", createInfo.shaders.vertex_src.as_path().string());

    createInfo.shaders.vertex_src = get_shader(vertex_path, true, false);

    createInfo.shader_input.bindings.push_back({ 14, GFXBindingType::StorageBuffer });

    if(positions_only) {
        createInfo.vertex_input.inputs = {
            {position_buffer_index, sizeof(prism::float3)},
            {bone_buffer_index, sizeof(BoneVertexData)}
        };

        createInfo.vertex_input.attributes = {
            {position_buffer_index, 0, 0, GFXVertexFormat::FLOAT3},
            {bone_buffer_index, 4, offsetof(BoneVertexData, ids), GFXVertexFormat::INT4},
            {bone_buffer_index, 5, offsetof(BoneVertexData, weights), GFXVertexFormat::FLOAT4}
        };
    } else {
        createInfo.vertex_input.inputs = {
            {position_buffer_index, sizeof(prism::float3)},
            {normal_buffer_index, sizeof(prism::float3)},
            {texcoord_buffer_index, sizeof(prism::float2)},
            {tangent_buffer_index, sizeof(prism::float3)},
            {bitangent_buffer_index, sizeof(prism::float3)},
            {bone_buffer_index, sizeof(BoneVertexData)}
        };

        createInfo.vertex_input.attributes = {
            {position_buffer_index, 0, 0, GFXVertexFormat::FLOAT3},
            {normal_buffer_index, 1, 0, GFXVertexFormat::FLOAT3},
            {texcoord_buffer_index, 2, 0, GFXVertexFormat::FLOAT2},
            {tangent_buffer_index, 3, 0, GFXVertexFormat::FLOAT3},
            {bitangent_buffer_index, 4, 0, GFXVertexFormat::FLOAT3},
            {bone_buffer_index, 5, offsetof(BoneVertexData, ids), GFXVertexFormat::INT4},
            {bone_buffer_index, 6, offsetof(BoneVertexData, weights), GFXVertexFormat::FLOAT4},
        };
    }

    return engine->get_gfx()->create_graphics_pipeline(createInfo);
}

std::tuple<GFXPipeline*, GFXPipeline*> MaterialCompiler::create_pipeline_permutations(GFXGraphicsPipelineCreateInfo& createInfo, bool positions_only) {
    auto st = create_static_pipeline(createInfo, positions_only);
    auto ss = create_skinned_pipeline(createInfo, positions_only);

    return {st, ss};
}

constexpr std::string_view struct_info =
R"(layout (constant_id = 0) const int max_materials = 25;
layout (constant_id = 1) const int max_lights = 25;
layout (constant_id = 2) const int max_spot_lights = 4;
layout (constant_id = 3) const int max_probes = 4;
struct Material {{
    vec4 color, info;
}};
struct Light {{
    vec4 positionType;
    vec4 directionPower;
    vec4 colorSize;
    vec4 shadowsEnable;
}};
struct Probe {{
    vec4 position, size;
}};
layout(std430, binding = 1) buffer readonly SceneInformation {{
    vec4 options;
    vec4 camPos;
    mat4 vp, lightSpace;
    mat4 spotLightSpaces[max_spot_lights];
    Material materials[max_materials];
    Light lights[max_lights];
    Probe probes[max_probes];
    int numLights;
}} scene;
layout (binding = 2) uniform sampler2D sun_shadow;
layout (binding = 6) uniform sampler2DArray spot_shadow;
layout(push_constant) uniform PushConstant {{
    mat4 model;
}};
)";

ShaderSource MaterialCompiler::compile_material_fragment(Material& material, bool use_ibl) {
    if(!render_options.enable_ibl)
        use_ibl = false;

    auto src = fmt::memory_buffer();

    switch(render_options.shadow_filter) {
        case ShadowFilter::None:
            format_to(std::back_inserter(src), "#define SHADOW_FILTER_NONE\n");
            break;
        case ShadowFilter::PCF:
            format_to(std::back_inserter(src), "#define SHADOW_FILTER_PCF\n");
            break;
        case ShadowFilter::PCSS:
            format_to(std::back_inserter(src), "#define SHADOW_FILTER_PCSS\n");
            break;
    }

    format_to(std::back_inserter(src),
        R"(layout (location = 0) in vec3 in_frag_pos;
        layout(location = 1) in vec3 in_normal;
        layout(location = 2) in vec2 in_uv;
        layout(location = 0) out vec4 frag_output;
        )");

    if(render_options.enable_point_shadows) {
        format_to(std::back_inserter(src),
                  R"(#define POINT_SHADOWS_SUPPORTED
                  layout (binding = 3) uniform samplerCubeArray point_shadow;
                  )");
    }

    format_to(std::back_inserter(src), struct_info);

    if(use_ibl) {
        format_to(std::back_inserter(src),
            R"(layout (binding = 7) uniform samplerCubeArray irrandianceSampler;
            layout (binding = 8) uniform samplerCubeArray prefilterSampler;
            layout (binding = 9) uniform sampler2D brdfSampler;
            )");
    }

    format_to(std::back_inserter(src),
        R"(layout(location = 4) in vec4 fragPosLightSpace;
        layout(location = 5) in mat3 in_tbn;
        layout(location = 14) in vec4 fragPostSpotLightSpace[max_spot_lights];
        layout(location = 13) in flat int inMaterialIndex;
        #include "common.glsl"
        #include "rendering.glsl"
        )");

    material.bound_textures.clear();

    // insert samplers as needed
    int sampler_index = 10;
    if(material.colorProperty.type == DataType::AssetTexture) {
        material.bound_textures[sampler_index] = material.colorProperty.value_tex;

        format_to(std::back_inserter(src), "layout(binding = {}) uniform sampler2D colorTexture;\n", sampler_index++);
    }

    if(material.normalProperty.type == DataType::AssetTexture) {
        material.bound_textures[sampler_index] = material.normalProperty.value_tex;

        format_to(std::back_inserter(src), "layout(binding = {}) uniform sampler2D normalTexture;\n", sampler_index++);
    }

    if(use_ibl) {
        format_to(std::back_inserter(src),
            R"(vec3 get_reflect(int i, vec3 final_normal) {{
            const vec3 direction = normalize(in_frag_pos - scene.camPos.xyz);
            const vec3 reflection = reflect(direction, normalize(final_normal));
            vec3 box_max = scene.probes[i].position.xyz + (scene.probes[i].size.xyz / 2.0f);
            vec3 box_min = scene.probes[i].position.xyz + -(scene.probes[i].size.xyz / 2.0f);
            vec3 unitary = vec3(1.0);
            vec3 first_plane_intersect = (box_max - in_frag_pos) / reflection;
            vec3 second_plane_intersect = (box_min - in_frag_pos) / reflection;
            vec3 furthest_plane = max(first_plane_intersect, second_plane_intersect);
            float distance = min(furthest_plane.x, min(furthest_plane.y, furthest_plane.z));
            vec3 intersect_position_world = in_frag_pos + reflection * distance;
            return intersect_position_world - scene.probes[i].position.xyz;
            }}
            )");
    }

    if(render_options.enable_normal_shadowing) {
        format_to(std::back_inserter(src),
            R"(float calculate_normal_lighting(in sampler2D normal_map, const vec3 normal, const vec3 light_dir) {{
            float height_scale = 0.8;
            float sample_count = 100.0;
            float inv_sample_count = 1.0 / sample_count;
            float hardness = 50 * 0.5;
            float lighting = clamp(dot(light_dir, normal), 0.0, 1.0);
            float slope = -lighting;
            vec2 dir = light_dir.xy * vec2(1.0, -1.0) * height_scale;
            float max_slope = 0.0;
            float step = inv_sample_count;
            float pos = step;
            pos = (-lighting >= 0.0) ? 1.001 : pos;
            vec2 noise = fract(in_frag_pos.xy * 0.5);
            noise.x = noise.x + noise.y * 0.5;
            pos = step - step * noise.x;
            float shadow = 0.0;
            while(pos <= 1.0) {{
                vec3 tmp_normal = texture(normal_map, in_uv + dir * pos).rgb;
                tmp_normal = in_tbn * (tmp_normal * 2.0 - 1.0);
                float tmp_lighting = dot(light_dir, tmp_normal);
                float shadowed = -tmp_lighting;
                slope += shadowed;
                if(slope > max_slope) {{
                    shadow += hardness * (1.0 - pos);
                }}
                max_slope = max(max_slope, slope);
                pos += step;
            }}
            return clamp(1.0 - shadow * inv_sample_count, 0.0, 1.0);
            }}
            )");
    }

    if(use_ibl) {
        format_to(std::back_inserter(src),
            R"(vec3 ibl(const int probe, const ComputedSurfaceInfo surface_info, const float intensity) {{
            const vec3 F = fresnel_schlick_roughness(surface_info.NdotV, surface_info.F0, surface_info.roughness);
            const vec3 R = get_reflect(probe, surface_info.N);
            const vec2 brdf = texture(brdfSampler, vec2(surface_info.NdotV, surface_info.roughness)).rg;
            const vec3 sampledIrradiance = texture(irrandianceSampler, vec4(surface_info.N, probe)).xyz;
            const vec3 prefilteredColor = textureLod(prefilterSampler, vec4(R, probe), surface_info.roughness * 4).xyz;
            const vec3 diffuse = sampledIrradiance * surface_info.diffuse_color;
            const vec3 specular = prefilteredColor * (F * brdf.x + brdf.y);
            return (diffuse + specular) * intensity;
            }}
            )");
    }

    format_to(std::back_inserter(src),  "void main() {{\n");

    if(material.colorProperty.type == DataType::Vector3) {
        format_to(std::back_inserter(src), "vec3 Color = vec3({}, {}, {});\n", material.colorProperty.value.x, material.colorProperty.value.y, material.colorProperty.value.z);
    } else if(material.colorProperty.type == DataType::AssetTexture) {
        format_to(std::back_inserter(src), "vec3 Color = texture(colorTexture, in_uv).rgb;\n");
    } else if(material.colorProperty.type == DataType::Float) {
        format_to(std::back_inserter(src),"vec3 Color = vec3({}});\n", material.colorProperty.float_value);
    }

    format_to(std::back_inserter(src),
              R"(vec3 final_diffuse_color = from_srgb_to_linear(Color);
              float final_roughness = scene.materials[inMaterialIndex].info.y;
              float final_metallic = scene.materials[inMaterialIndex].info.x;
              )");

    if(material.normalProperty.type == DataType::AssetTexture) {
        format_to(std::back_inserter(src),
                  R"(vec3 final_normal = texture(normalTexture, in_uv).rgb;
                  final_normal = final_normal * 2.0 - 1.0;
                  final_normal = in_tbn * final_normal;
                  )");
    } else {
        format_to(std::back_inserter(src), "vec3 final_normal = in_normal;\n");
    }

    format_to(std::back_inserter(src),
    R"(ComputedSurfaceInfo surface_info = compute_surface(final_diffuse_color.rgb, final_normal, final_metallic, final_roughness);
    vec3 Lo = vec3(0);
    for(int i = 0; i < scene.numLights; i++) {{
        const int type = int(scene.lights[i].positionType.w);
        ComputedLightInformation light_info;
        switch(type) {{
            case 0:
                light_info = calculate_point(scene.lights[i]);
                break;
            case 1:
                light_info = calculate_spot(scene.lights[i]);
                break;
            case 2:
                light_info = calculate_sun(scene.lights[i]);
                break;
        }}
    SurfaceBRDF surface_brdf = brdf(light_info.direction, surface_info);
    )");

    if(render_options.enable_normal_mapping && material.normalProperty.type == DataType::AssetTexture && render_options.enable_normal_shadowing) {
        format_to(std::back_inserter(src), "light_info.radiance *= calculate_normal_lighting(normalTexture, final_normal, light_info.direction);\n");
    }

    format_to(std::back_inserter(src), "Lo += ((surface_brdf.specular + surface_brdf.diffuse) * light_info.radiance * surface_brdf.NdotL) * scene.lights[i].colorSize.rgb;\n \
    }}\n");

    if(use_ibl) {
        format_to(std::back_inserter(src),
        R"(vec3 ambient = vec3(0.0);
        float sum = 0.0;
        for(int i = 0; i < max_probes; i++) {{
            if(scene.probes[i].position.w == 1) {{
                const vec3 position = scene.probes[i].position.xyz;
                const vec3 probe_min = position - (scene.probes[i].size.xyz / 2.0);
                const vec3 probe_max = position + (scene.probes[i].size.xyz / 2.0);
                if(all(greaterThan(in_frag_pos, probe_min)) && all(lessThan(in_frag_pos, probe_max))) {{
                    float intensity = 1.0 - length(abs(in_frag_pos - position) / (scene.probes[i].size.xyz / 2.0));
                    intensity = clamp(intensity, 0.0, 1.0) * scene.probes[i].size.w;
                    ambient += ibl(i, surface_info, intensity);
                    sum += intensity;
                }}
            }} else if(scene.probes[i].position.w == 2) {{
                ambient += ibl(i, surface_info, scene.probes[i].size.w);
                sum += scene.probes[i].size.w;
            }}
        }}
        ambient /= sum;
        )");

        format_to(std::back_inserter(src), "frag_output = vec4(ambient + Lo, 1.0);\n");
    } else {
        format_to(std::back_inserter(src), "frag_output = vec4(Lo, 1.0);\n");
    }

    format_to(std::back_inserter(src), "}}\n");

    return *shader_compiler.compile(ShaderLanguage::GLSL, ShaderStage::Fragment, ShaderSource(to_string(src)), engine->get_gfx()->accepted_shader_language());
}