// This file is part of the 64k demo project.
// It implements the ShaderComposer class.

#include "gpu/effects/shader_composer.h"
#include <set>
#include <sstream>

ShaderComposer& ShaderComposer::Get() {
  static ShaderComposer instance;
  return instance;
}

void ShaderComposer::RegisterSnippet(const std::string& name,
                                     const std::string& code) {
  snippets_[name] = code;
}

void ShaderComposer::ResolveRecursive(const std::string& source,
                                      std::stringstream& ss,
                                      std::set<std::string>& included,
                                      const CompositionMap& substitutions) {
  std::istringstream stream(source);
  std::string line;
  while (std::getline(stream, line)) {
    // Check for #include "snippet_name"
    if (line.compare(0, 9, "#include ") == 0) {
      size_t start = line.find('"');
      size_t end = line.find('"', start + 1);
      if (start != std::string::npos && end != std::string::npos) {
        std::string name = line.substr(start + 1, end - start - 1);

        // Apply substitution if available
        auto sub_it = substitutions.find(name);
        if (sub_it != substitutions.end()) {
          name = sub_it->second;
        }

        if (included.find(name) == included.end()) {
          included.insert(name);
          auto it = snippets_.find(name);
          if (it != snippets_.end()) {
            ss << "// --- Included: " << name << " ---\n";
            ResolveRecursive(it->second, ss, included, substitutions);
            ss << "// --- End Include: " << name << " ---\n";
          } else {
            ss << "// ERROR: Snippet not found: " << name << "\n";
          }
        }
      }
    } else {
      ss << line << "\n";
    }
  }
}

std::string
ShaderComposer::Compose(const std::vector<std::string>& dependencies,
                        const std::string& main_code,
                        const CompositionMap& substitutions) {
  std::stringstream ss;
  ss << "// Generated by ShaderComposer\n\n";

  std::set<std::string> included;

  // Process explicit dependencies first
  for (const auto& dep : dependencies) {
    std::string name = dep;
    auto sub_it = substitutions.find(name);
    if (sub_it != substitutions.end()) {
      name = sub_it->second;
    }

    if (included.find(name) == included.end()) {
      included.insert(name);
      auto it = snippets_.find(name);
      if (it != snippets_.end()) {
        ss << "// --- Dependency: " << name << " ---\n";
        ResolveRecursive(it->second, ss, included, substitutions);
        ss << "\n";
      }
    }
  }

  ss << "// --- Main Code ---\n";
  ResolveRecursive(main_code, ss, included, substitutions);

  return ss.str();
}

void ShaderComposer::VerifyIncludes() const {
#if !defined(STRIP_ALL)
  // Known placeholders that get substituted at composition time
  std::set<std::string> known_placeholders = {"render/scene_query_mode"};

  std::set<std::string> missing;
  for (const auto& [name, code] : snippets_) {
    std::istringstream stream(code);
    std::string line;
    while (std::getline(stream, line)) {
      if (line.compare(0, 9, "#include ") == 0) {
        size_t start = line.find('"');
        size_t end = line.find('"', start + 1);
        if (start != std::string::npos && end != std::string::npos) {
          std::string included = line.substr(start + 1, end - start - 1);
          if (snippets_.find(included) == snippets_.end() &&
              known_placeholders.find(included) == known_placeholders.end()) {
            missing.insert(included);
          }
        }
      }
    }
  }
  if (!missing.empty()) {
    fprintf(stderr, "WARNING: Unregistered shader snippets:\n");
    for (const auto& name : missing) {
      fprintf(stderr, "  - %s\n", name.c_str());
    }
  }
#endif
}