#include <filesystem>
#include <sstream>

#include "log.hpp"
#include "shadercompiler.hpp"
#include "string_utils.hpp"

bool has_extension(const std::filesystem::path& path, const std::string_view extension) {
    return string_contains(path.filename().string(), extension);
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        prism::log("Not enough arguments!");
        return -1;
    }

    shader_compiler.set_include_path(std::filesystem::current_path().string());

    std::filesystem::path source_path = argv[1];
    std::filesystem::path destination_path = argv[2];

    ShaderLanguage language = ShaderLanguage::SPIRV;
    std::string_view shader_language_string = argv[3];
    if (shader_language_string == "spv") {
        language = ShaderLanguage::SPIRV;
    } else if (shader_language_string == "msl") {
        language = ShaderLanguage::MSL;
    } else if (shader_language_string == "glsl") {
        language = ShaderLanguage::GLSL;
    } else if (shader_language_string == "hlsl") {
        language = ShaderLanguage::HLSL;
    }

    CompileOptions options;

    std::string_view extra_options = argc > 4 ? argv[4] : "";
    if (extra_options == "mobile")
        options.is_apple_mobile = true;

    std::ifstream t(source_path);
    std::stringstream buffer;
    buffer << t.rdbuf();

    if (has_extension(source_path, "nocompile")) {
        destination_path = remove_substring(destination_path.string(), ".nocompile"); // remove extension

        std::ofstream out(destination_path);
        out << buffer.rdbuf();
    } else {
        ShaderStage stage = ShaderStage::Vertex;
        if (has_extension(source_path, ".vert"))
            stage = ShaderStage::Vertex;
        else if (has_extension(source_path, ".frag"))
            stage = ShaderStage::Fragment;
        else if (has_extension(source_path, ".comp"))
            stage = ShaderStage::Compute;

        const auto compiled_source =
            shader_compiler.compile(ShaderLanguage::GLSL, stage, ShaderSource(buffer.str()), language, options);
        if (!compiled_source.has_value()) {
            prism::log("Error when compiling {}!", source_path.string());
            return -1;
        }

        switch (language) {
            // right now, WGSL is outputted as SPIR-V with some WGSL compatibility stuff included
            case ShaderLanguage::SPIRV: {
                const auto spirv = compiled_source->as_bytecode();

                std::ofstream out(destination_path, std::ios::binary); // remove .glsl
                out.write((char*)spirv.data(), spirv.size() * sizeof(uint32_t));
            } break;
            case ShaderLanguage::MSL: {
                std::ofstream out(destination_path); // remove .glsl
                out << compiled_source->as_string();
            } break;
            case ShaderLanguage::HLSL: {
                std::ofstream out(destination_path); // remove .glsl
                out << compiled_source->as_string();
            } break;
            default:
                break;
        }
    }

    // TODO: output disabled for now, will have to expose in a better arg system
    // prism::log::info(System::Core, "Successfully written shader from {} to {}.", source_path, destination_path);

    return 0;
}