#pragma once

#include <filesystem>
#include <fmt/base.h>
#include <fstream>
#include <shaderc/shaderc.hpp>
#include <string>
#include <vector>

#include "types.hpp"

class ShaderCompiler {
 public:
  ShaderCompiler() = default;

  /**
   * @brief Compiles or loads a cached SPIR-V shader from a file.
   *
   * @param shaderSource The source code of the shader.
   * @param kind The type of shader being compiled (vertex, fragment, etc.).
   * @param shaderName The name used for caching the compiled shader.
   * @return std::vector<u32> A vector containing the compiled SPIR-V code.
   * @throws std::runtime_error if shader compilation fails.
   *
   * This function attempts to load a shader from the cache. If the shader
   * is not found in the cache, it compiles the shader from the source code
   * and saves the result to the cache for future use.
   */
  static fn getCompiledShader(
    const char*                shaderSource,
    const shaderc_shader_kind& kind,
    const string&              shaderName
  ) -> std::vector<u32> {
    using namespace std;

    const string cacheFile = getCacheFilePath(shaderName);

    // Try loading from cache first
    vector<u32> spirvCode = loadCachedShader(cacheFile);

    if (!spirvCode.empty()) {
      fmt::println("Loaded shader from cache: {}", cacheFile);
      return spirvCode;
    }

    // Cache miss, compile the shader
    spirvCode = compileShader(shaderSource, kind);

    if (spirvCode.empty())
      throw runtime_error("Shader compilation failed for: " + shaderName);

    // Cache the compiled SPIR-V binary
    saveCompiledShader(spirvCode, cacheFile);
    return spirvCode;
  }

 private:
  /**
   * @brief Generates the cache file path based on the operating system.
   *
   * @param shaderName The name used for the shader file.
   * @return string The full path to the cache file.
   *
   * This function determines the appropriate directory for caching shaders
   * based on the operating system and returns the full path for the specified shader name.
   */
  static fn getCacheFilePath(const string& shaderName) -> string {
    using namespace std::filesystem;

#ifdef _WIN32
    path cacheDir = path(getenv("LOCALAPPDATA")) / "VulkanApp" / "Shaders";
#elif defined(__APPLE__)
    path cacheDir = path(getenv("HOME")) / "Library" / "Application Support" / "VulkanApp" / "Shaders";
#else // Assume Linux or other UNIX-like systems
    path cacheDir = path(getenv("HOME")) / ".local" / "share" / "VulkanApp" / "Shaders";
#endif

    if (!exists(cacheDir))
      create_directories(cacheDir); // Create the directory if it doesn't exist

    return (cacheDir / (shaderName + ".spv")).string();
  }

  /**
   * @brief Compiles GLSL code to SPIR-V.
   *
   * @param source The GLSL source code to compile.
   * @param kind The type of shader being compiled.
   * @return std::vector<u32> A vector containing the compiled SPIR-V code.
   *
   * This function uses the shaderc library to compile GLSL source code into
   * SPIR-V binary format. If the compilation fails, an empty vector is returned.
   */
  static fn compileShader(const char* source, const shaderc_shader_kind& kind) -> std::vector<u32> {
    using namespace shaderc;

    Compiler       compiler;
    CompileOptions options;

    SpvCompilationResult result = compiler.CompileGlslToSpv(source, kind, "shader.glsl", "main", options);

    if (result.GetCompilationStatus() != shaderc_compilation_status_success) {
      fmt::println(stderr, "Shader compilation failed: {}", result.GetErrorMessage());
      return {};
    }

    return { result.cbegin(), result.cend() };
  }

  /**
   * @brief Loads compiled SPIR-V shader code from a cache file.
   *
   * @param cacheFile The path to the cached SPIR-V file.
   * @return std::vector<u32> A vector containing the loaded SPIR-V code.
   * @throws std::runtime_error if the file cannot be read.
   *
   * This function checks if the specified cache file exists and reads its
   * contents into a vector. If the file cannot be opened, an exception is thrown.
   */
  static fn loadCachedShader(const string& cacheFile) -> std::vector<u32> {
    using namespace std;

    if (!filesystem::exists(cacheFile))
      return {}; // No cached file

    ifstream file(cacheFile, ios::binary);

    // Check if the file was successfully opened
    if (!file)
      throw runtime_error("Failed to open cached shader file: " + cacheFile);

    usize       fileSize = filesystem::file_size(cacheFile);
    vector<u32> spirvCode(fileSize / sizeof(u32));

    // Read entire file content into the vector
    if (!file.read(bit_cast<char*>(spirvCode.data()), static_cast<streamsize>(fileSize)))
      throw runtime_error("Failed to read cached shader file: " + cacheFile);

    return spirvCode;
  }

  /**
   * @brief Saves compiled SPIR-V binary to a cache file.
   *
   * @param spirvCode The SPIR-V code to save.
   * @param cacheFile The path to the file where the SPIR-V code will be saved.
   * @throws std::runtime_error if the file cannot be written.
   *
   * This function writes the compiled SPIR-V binary to the specified file.
   * If the file cannot be opened or written, an exception is thrown.
   */
  static fn saveCompiledShader(const std::vector<u32>& spirvCode, const string& cacheFile) -> void {
    using namespace std;

    ofstream file(cacheFile, ios::binary);

    // Check if the file was successfully opened
    if (!file)
      throw runtime_error("Failed to open file for saving shader: " + cacheFile);

    if (!file.write(
          bit_cast<const char*>(spirvCode.data()), static_cast<streamsize>(spirvCode.size() * sizeof(u32))
        ))
      throw runtime_error("Failed to save shader to cache: " + cacheFile);
  }
};
