#include "glsl.hpp"
#include <glslang/Public/ShaderLang.h>

using namespace std;

namespace shader_manager
{
    namespace glsl
    {
        // utils

        template<typename T>
        vector<T> create_reserved_vec(size_t size)
        {
            vector<T> result;
            result.reserve(size);
            return result;
        }

        // converting

        glslang::EShSource to_glslang_source_language(source_language language)
        {
            switch (language)
            {
            case source_language::glsl_450:
                return ::glslang::EShSourceGlsl;

            case source_language::glsl_400:
                return ::glslang::EShSourceGlsl;

            case source_language::glsl_330:
                return ::glslang::EShSourceGlsl;

            case source_language::hlsl_500:
                return ::glslang::EShSourceHlsl;

            case source_language::hlsl_400:
                return ::glslang::EShSourceHlsl;

            case source_language::hlsl_300:
                return ::glslang::EShSourceHlsl;

            default:
                throw runtime_error("undefined enum value");
            }
        }

        unsigned to_glslang_graphic_api_version(source_language language)
        {
            switch (language)
            {
            case source_language::glsl_450:
                return 450;

            case source_language::glsl_400:
                return 400;

            case source_language::glsl_330:
                return 330;

            case source_language::hlsl_500:
                return 500;

            case source_language::hlsl_400:
                return 400;

            case source_language::hlsl_300:
                return 300;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        glslang::EShTargetLanguage to_glslang_target_language(target_language language)
        {
            switch (language)
            {
            case target_language::spv_1_0:
                return ::glslang::EShTargetSpv;

            case target_language::spv_1_1:
                return ::glslang::EShTargetSpv;

            case target_language::spv_1_2:
                return ::glslang::EShTargetSpv;

            case target_language::spv_1_3:
                return ::glslang::EShTargetSpv;

            case target_language::spv_1_4:
                return ::glslang::EShTargetSpv;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        glslang::EShTargetLanguageVersion to_glslang_target_language_version(target_language language)
        {
            switch (language)
            {
            case target_language::spv_1_0:
                return ::glslang::EShTargetSpv_1_0;

            case target_language::spv_1_1:
                return ::glslang::EShTargetSpv_1_1;

            case target_language::spv_1_2:
                return ::glslang::EShTargetSpv_1_2;

            case target_language::spv_1_3:
                return ::glslang::EShTargetSpv_1_3;

            case target_language::spv_1_4:
                return ::glslang::EShTargetSpv_1_4;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        glslang::EShClient to_glslang_graphic_api(graphic_api api)
        {
            switch (api)
            {
            case graphic_api::openGL_450:
                return ::glslang::EShClientOpenGL;

            case graphic_api::vulkan_1_0:
                return ::glslang::EShClientVulkan;

            case graphic_api::vulkan_1_1:
                return ::glslang::EShClientVulkan;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        glslang::EShTargetClientVersion to_glslang_graphic_api_version(graphic_api api)
        {
            switch (api)
            {
            case graphic_api::openGL_450:
                return ::glslang::EShTargetOpenGL_450;

            case graphic_api::vulkan_1_0:
                return ::glslang::EShTargetVulkan_1_0;

            case graphic_api::vulkan_1_1:
                return ::glslang::EShTargetVulkan_1_1;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        EShLanguage to_glslang_shader_stage(shader_stage stage)
        {
            switch (stage)
            {
            case shader_stage::compute:
                return EShLangCompute;

            case shader_stage::fragment:
                return EShLangFragment;

            case shader_stage::geometry:
                return EShLangGeometry;

            case shader_stage::tessellation_control:
                return EShLangTessControl;

            case shader_stage::tessellation_evaluation:
                return EShLangTessEvaluation;

            case shader_stage::vertex:
                return EShLangVertex;

            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        vector<char const*> to_cstr_vec(vector<string> const& vec)
        {
            std::vector<char const*> result;
            for (auto const& str : vec)
            {
                result.push_back(str.c_str());
            }
            return result;
        }

        // initialization of program

        void set_shader_options(glslang::TShader& shader, shader_options options)
        {
            switch (options)
            {
            case shader_options::auto_map_bindings:
                shader.setAutoMapBindings(true);
                break;
            case shader_options::auto_map_locations:
                shader.setAutoMapLocations(true);
                break;
            case shader_options::invert_y:
                shader.setInvertY(true);
                break;
            case shader_options::hlsl_io_mapping:
                shader.setHlslIoMapping(true);
                break;
            case shader_options::flatten_uniform_arrays:
                shader.setFlattenUniformArrays(true);
                break;
            case shader_options::no_storage_format:
                shader.setNoStorageFormat(true);
                break;
            default:
                throw std::runtime_error("undefined enum value");
            }
        }

        void configurate_shader(glslang::TShader& gl_shader, shader shader)
        {
            gl_shader.setEntryPoint(shader.entry_point.data());
            gl_shader.setEnvClient(
                to_glslang_graphic_api(shader.graphic_api),
                to_glslang_graphic_api_version(shader.graphic_api)
            );
            gl_shader.setEnvInput(
                to_glslang_source_language(shader.source_language),
                to_glslang_shader_stage(shader.stage),
                to_glslang_graphic_api(shader.graphic_api),
                to_glslang_graphic_api_version(shader.graphic_api)
            );
        }

        void prepare_shader(glslang::TShader& glang_shader, vector<char const*> const& code,
            shader const& shader)
        {
            static TBuiltInResource resource;

            glang_shader.setStrings(code.data(), static_cast<int>(code.size()));
            configurate_shader(glang_shader, shader);
            glang_shader.parse(&resource, 0, false, EShMessages::EShMsgDefault);
        }

        void init_program(glslang::TProgram& program, std::vector<shader> const& shaders)
        {
            auto glang_shaders = create_reserved_vec<glslang::TShader>(shaders.size());
            auto shaders_code = create_reserved_vec<vector<char const*>>(shaders.size());

            for (auto const& shader : shaders)
            {
                shaders_code.push_back(to_cstr_vec(shader.source));

                glslang::TShader glang_shader(to_glslang_shader_stage(shader.stage));
                prepare_shader(glang_shader, shaders_code.back(), shader);
                glang_shaders.push_back(glang_shader);

                program.addShader(&glang_shader);
            }

            if (!program.link(EShMsgDefault))
            {
                throw std::runtime_error("fail to link shaders");
            }
            if (!program.mapIO())
            {
                throw std::runtime_error("fail to map IO");
            }
            if (!program.buildReflection())
            {
                throw std::runtime_error("fail to build_reflection");
            }
        }

        // variable factory

        class glslang_variable : public variable, public opengl_fundamental_variable_impl
        {
        public:
            glslang_variable(glslang::TObjectReflection const& var_info) :
                opengl_fundamental_variable_impl(
                    static_cast<opengl_fundamental_type>(var_info.glDefineType)), name_(var_info.name)
            {
            }

            std::string get_name() const override
            {
                return name_;
            }

        private:
            std::string name_;
        };

        shared_ptr<variable> create_variable(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info, bool in_struct);

        vector<shared_ptr<variable>> create_variables(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info)
        {
            vector<shared_ptr<variable>> variables;
            for (utils::index_t i = var_info.arrayStride; i != var_info.topLevelArrayStride; ++i)
            {
                variables.push_back(create_variable(program, program.getBufferVariable(static_cast<int>(i)), true));
            }
            return variables;
        }

        class struct_glslang_variable : public variable, public variable_impl, public struct_variable_impl
        {
        public:
            struct_glslang_variable(glslang::TProgram const& program, glslang::TObjectReflection const& var_info) :
                variable_impl(var_info.name, "struct"),
                struct_variable_impl(create_variables(program, var_info))
            {
            }
        };

        shared_ptr<variable> create_array_in_struct_variable(glslang::TObjectReflection const& var_info)
        {
            struct var : variable, glslang_variable, in_struct_variable_impl, array_variable_impl
            {
                var(glslang::TObjectReflection const& var_info) :
                    glslang_variable(var_info),
                    in_struct_variable_impl(var_info.offset),
                    array_variable_impl(var_info.size)
                {
                }
            };

            return make_shared<var>(var_info);
        }

        shared_ptr<variable> create_array_variable(glslang::TObjectReflection const& var_info)
        {
            struct var : variable, glslang_variable, array_variable_impl
            {
                var(glslang::TObjectReflection const& var_info) :
                    glslang_variable(var_info),
                    array_variable_impl(var_info.size)
                {
                }
            };

            return make_shared<var>(var_info);
        }

        shared_ptr<variable> create_in_struct_variable(glslang::TObjectReflection const& var_info)
        {
            struct var : variable, glslang_variable, in_struct_variable_impl
            {
                var(glslang::TObjectReflection const& var_info) :
                    glslang_variable(var_info),
                    in_struct_variable_impl(var_info.offset)
                {
                }
            };

            return make_shared<var>(var_info);
        }

        shared_ptr<variable> create_struct_array_in_struct_variable(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info)
        {
            struct var : variable, struct_glslang_variable, in_struct_variable_impl, array_variable_impl
            {
                var(glslang::TProgram& program, glslang::TObjectReflection const& var_info) :
                    struct_glslang_variable(program, var_info),
                    in_struct_variable_impl(var_info.offset),
                    array_variable_impl(var_info.size)
                {
                }
            };

            return make_shared<var>(program, var_info);
        }

        shared_ptr<variable> create_struct_array_variable(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info)
        {
            struct var : variable, struct_glslang_variable, array_variable_impl
            {
                var(glslang::TProgram& program, glslang::TObjectReflection const& var_info) :
                    struct_glslang_variable(program, var_info),
                    array_variable_impl(var_info.size)
                {
                }
            };

            return make_shared<var>(program, var_info);
        }

        shared_ptr<variable> create_struct_in_struct_variable(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info)
        {
            struct var : variable, struct_glslang_variable, in_struct_variable_impl
            {
                var(glslang::TProgram& program, glslang::TObjectReflection const& var_info) :
                    struct_glslang_variable(program, var_info),
                    in_struct_variable_impl(var_info.offset)
                {
                }
            };

            return make_shared<var>(program, var_info);
        }

        // shader reflection

        shared_ptr<variable> create_variable(glslang::TProgram const& program,
            glslang::TObjectReflection const& var_info, bool in_struct)
        {
            if (var_info.numMembers)
            {
                if (var_info.size && in_struct)
                {
                    return create_struct_array_in_struct_variable(program, var_info);
                }
                else if (var_info.size)
                {
                    return create_struct_array_variable(program, var_info);
                }
                else if (in_struct)
                {
                    return create_struct_in_struct_variable(program, var_info);
                }
                else
                {
                    return make_shared<struct_glslang_variable>(program, var_info);
                }
            }
            else
            {
                if (var_info.size && in_struct)
                {
                    return create_array_in_struct_variable(var_info);
                }
                else if (var_info.size)
                {
                    return create_array_variable(var_info);
                }
                else if (in_struct)
                {
                    return create_in_struct_variable(var_info);
                }
                else
                {
                    return make_shared<glslang_variable>(var_info);
                }
            }
        }

        void get_uniform_variables(glslang::TProgram const& program, vector<shared_ptr<resource>>& resources)
        {
            for (utils::index_t i = 0; i != program.getNumUniformVariables(); ++i)
            {
                auto const& var_info = program.getUniform(static_cast<int>(i));
                resources.push_back(make_shared<uniform_resource_impl>(
                    create_variable(program, var_info, false), "uniform", var_info.getBinding()));
            }

            for (utils::index_t i = 0; i != program.getNumUniformBlocks(); ++i)
            {
                auto const& var_info = program.getUniformBlock(static_cast<int>(i));
                resources.push_back(make_shared<uniform_resource_impl>(
                    create_variable(program, var_info, false), "uniform", var_info.getBinding()));
            }
        }

        void get_input_variables(glslang::TProgram const& program, vector<shared_ptr<resource>>& resources)
        {
            for (utils::index_t i = 0; i != program.getNumPipeInputs(); ++i)
            {
                auto const& var_info = program.getPipeInput(static_cast<int>(i));
                resources.push_back(make_shared<input_resource_impl>(
                    create_variable(program, var_info, false), "input", var_info.getBinding()));
            }
        }

        void get_output_variables(glslang::TProgram const& program, vector<shared_ptr<resource>>& resources)
        {
            for (utils::index_t i = 0; i != program.getNumPipeInputs(); ++i)
            {
                auto const& var_info = program.getPipeOutput(static_cast<int>(i));
                resources.push_back(make_shared<output_resource_impl>(
                    create_variable(program, var_info, false), "output", var_info.getBinding()));
            }
        }

        vector<shared_ptr<resource>> build_resources(glslang::TProgram const& program)
        {
            vector<shared_ptr<resource>> result;
            get_uniform_variables(program, result);
            get_input_variables(program, result);
            get_output_variables(program, result);

            return result;
        }

        // analizer ---------------------------------------------------------------------------------------------

        std::vector<std::shared_ptr<resource>> analize_shaders(std::vector<shader> const& shaders)
        {
            glslang::InitializeProcess();
            glslang::TProgram program;
            init_program(program, shaders);
            auto resources = build_resources(program);
            glslang::FinalizeProcess();
            return resources;
        }
    }
}