#ifndef SHADERCOMPILER_H
#define SHADERCOMPILER_H

#include <cctype>
#include <filesystem>
#include <fstream>
#include <regex>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_set>
#include <vector>

// glslang
#include <glslang/Public/ResourceLimits.h>
#include <glslang/Public/ShaderLang.h>
#include <SPIRV/GlslangToSpv.h>
#include <StandAlone/DirStackFileIncluder.h>

namespace fs = std::filesystem;

class SpirvShaderCompiler
{
public:
    struct Options
    {
        std::vector<fs::path> includeDirs;        // searched for #include "..."
        std::vector<std::string> defines;         // "NAME" or "NAME=VALUE"
        bool debugInfo = false;                   // emit debug info in SPIR-V
        glslang::EShTargetClientVersion vulkanTarget = glslang::EShTargetVulkan_1_2;
        glslang::EShTargetLanguageVersion spirvTarget = glslang::EShTargetSpv_1_5;
        int glslVersion = 450;
    };

    explicit SpirvShaderCompiler(Options opt = {})
        : m_opt(std::move(opt))
    {
        ensureInitialized();
    }

    // Returns true if rebuilt. Throws on compilation/dependency errors.
    bool CompileFileIfNeeded(const fs::path& entryFile, const fs::path& outSpv)
    {
        if (!IsEntryShader(entryFile))
            throw std::runtime_error("CompileFileIfNeeded called with non-entry shader: " + entryFile.string());

        if (!NeedsRebuild(entryFile, outSpv))
            return false;

        const EShLanguage stage = StageFromExtension(entryFile);

        auto spirv = CompileToSpirvWords(entryFile, stage);
        WriteSpv(outSpv, spirv);
        return true;
    }

    // Utility for your cooker rules
    static bool IsEntryShader(const fs::path& p)
    {
        const auto ext = ToLower(p.extension().string());
        return ext == ".vert" || ext == ".frag" || ext == ".comp";
    }

    static bool IsIncludeOnlyShader(const fs::path& p)
    {
        return ToLower(p.extension().string()) == ".glsl";
    }

private:
    Options m_opt;

    // --- init ---
    static void ensureInitialized()
    {
        static bool initialized = false;
        if (!initialized) {
            glslang::InitializeProcess();
            initialized = true;
        }
    }

    // --- file helpers ---
    static std::string ReadTextFile(const fs::path& p)
    {
        std::ifstream f(p, std::ios::in);
        if (!f) throw std::runtime_error("Failed to open file: " + p.string());
        std::stringstream ss;
        ss << f.rdbuf();
        return ss.str();
    }

    static void EnsureParentDir(const fs::path& p)
    {
        fs::create_directories(p.parent_path());
    }

    static void WriteSpv(const fs::path& out, const std::vector<uint32_t>& spirv)
    {
        EnsureParentDir(out);
        std::ofstream f(out, std::ios::binary);
        if (!f) throw std::runtime_error("Failed to write: " + out.string());
        f.write(reinterpret_cast<const char*>(spirv.data()),
                static_cast<std::streamsize>(spirv.size() * sizeof(uint32_t)));
    }

    static std::string ToLower(std::string s)
    {
        for (char& c : s) c = (char)std::tolower((unsigned char)c);
        return s;
    }

    static EShLanguage StageFromExtension(const fs::path& p)
    {
        const auto ext = ToLower(p.extension().string());
        if (ext == ".vert") return EShLangVertex;
        if (ext == ".frag") return EShLangFragment;
        if (ext == ".comp") return EShLangCompute;
        throw std::runtime_error("Unknown entry shader stage extension: " + p.string());
    }

    static std::vector<std::string> FindIncludes(const std::string& text)
    {
        std::vector<std::string> out;
        static const std::regex re(
            R"(^\s*#\s*include\s*"([^"]+)\"\s*$)",
            std::regex::icase
        );

        std::string line;
        std::istringstream iss(text);
        while (std::getline(iss, line)) {
            // strip single-line comments
            auto pos = line.find("//");
            if (pos != std::string::npos) line = line.substr(0, pos);

            std::smatch m;
            if (std::regex_search(line, m, re)) {
                out.push_back(m[1].str());
            }
        }
        return out;
    }

    fs::path ResolveInclude(const fs::path& includingFile, const std::string& includeName) const
    {
        {
            fs::path p = includingFile.parent_path() / includeName;
            if (fs::exists(p)) return fs::weakly_canonical(p);
        }
        for (const auto& d : m_opt.includeDirs) {
            fs::path p = d / includeName;
            if (fs::exists(p)) return fs::weakly_canonical(p);
        }

        throw std::runtime_error(
            "Failed to resolve include \"" + includeName + "\" included from " + includingFile.string()
        );
    }

    void CollectDepsRecursive(const fs::path& file, std::unordered_set<std::string>& visitedAbs) const
    {
        fs::path abs = fs::weakly_canonical(file);
        const std::string key = abs.string();
        if (visitedAbs.contains(key)) return;
        visitedAbs.insert(key);

        const std::string text = ReadTextFile(abs);
        for (const auto& inc : FindIncludes(text)) {
            fs::path resolved = ResolveInclude(abs, inc);
            CollectDepsRecursive(resolved, visitedAbs);
        }
    }

    fs::file_time_type NewestDependencyTime(const fs::path& entryFile) const
    {
        std::unordered_set<std::string> deps;
        CollectDepsRecursive(entryFile, deps);

        fs::file_time_type newest = fs::file_time_type::min();
        for (const auto& dep : deps) {
            auto t = fs::last_write_time(dep);
            if (t > newest) newest = t;
        }
        return newest;
    }

    bool NeedsRebuild(const fs::path& entryFile, const fs::path& outSpv) const
    {
        if (!fs::exists(outSpv)) return true;

        const auto outTime = fs::last_write_time(outSpv);
        const auto newestDep = NewestDependencyTime(entryFile);
        return newestDep > outTime;
    }

    // --- compilation ---
    std::vector<uint32_t> CompileToSpirvWords(const fs::path& entryFile, EShLanguage stage) const
    {
        const std::string source = ReadTextFile(entryFile);

        std::string preamble;
        for (const auto& d : m_opt.defines)
            preamble += "#define " + d + "\n";

        const char* strings[] = { source.c_str() };

        glslang::TShader shader(stage);
        shader.setStrings(strings, 1);
        shader.setPreamble(preamble.c_str());
        shader.setEntryPoint("main");
        shader.setSourceEntryPoint("main");

        shader.setEnvInput(glslang::EShSourceGlsl, stage, glslang::EShClientVulkan, 100);
        shader.setEnvClient(glslang::EShClientVulkan, m_opt.vulkanTarget);
        shader.setEnvTarget(glslang::EShTargetSpv, m_opt.spirvTarget);

        DirStackFileIncluder includer;
        // Search local first, then include dirs:
        includer.pushExternalLocalDirectory(entryFile.parent_path().string());
        for (auto& d : m_opt.includeDirs)
            includer.pushExternalLocalDirectory(d.string());

        const TBuiltInResource* resources = GetDefaultResources();
        const EShMessages messages = (EShMessages)(EShMsgSpvRules | EShMsgVulkanRules);

        if (!shader.parse(resources, m_opt.glslVersion, false, messages, includer)) {
            throw std::runtime_error(
                "GLSL parse failed: " + entryFile.string() + "\n" +
                shader.getInfoLog() + "\n" + shader.getInfoDebugLog()
            );
        }

        glslang::TProgram program;
        program.addShader(&shader);

        if (!program.link(messages)) {
            throw std::runtime_error(
                "GLSL link failed: " + entryFile.string() + "\n" +
                program.getInfoLog() + "\n" + program.getInfoDebugLog()
            );
        }

        std::vector<uint32_t> spirv;
        glslang::SpvOptions spvOpt{};
        spvOpt.generateDebugInfo = m_opt.debugInfo;
        glslang::GlslangToSpv(*program.getIntermediate(stage), spirv, &spvOpt);
        return spirv;
    }
};

#endif //SHADERCOMPILER_H