// This file is part of the 64k demo project.
// It tests the ShaderComposer utility.

#include "gpu/effects/shader_composer.h"
#include <cassert>
#include <iostream>
#include <string>

#if defined(USE_TEST_ASSETS)
#include "test_assets.h"
#else
#include "generated/assets.h"
#endif

// Forward declaration for asset loading
const uint8_t* GetAsset(AssetId asset_id, size_t* out_size);

void test_composition() {
  std::cout << "Testing Shader Composition..." << std::endl;
  auto& sc = ShaderComposer::Get();

  sc.RegisterSnippet("math", "fn add(a: f32, b: f32) -> f32 { return a + b; }");
  sc.RegisterSnippet("util", "fn square(a: f32) -> f32 { return a * a; }");

  std::string main_code = "fn main() { let x = add(1.0, square(2.0)); }";
  std::string result = sc.Compose({"math", "util"}, main_code);

  // Verify order and presence
  assert(result.find("Snippet: math") != std::string::npos);
  assert(result.find("Snippet: util") != std::string::npos);
  assert(result.find("Main Code") != std::string::npos);

  size_t pos_math = result.find("Snippet: math");
  size_t pos_util = result.find("Snippet: util");
  size_t pos_main = result.find("Main Code");

  assert(pos_math < pos_util);
  assert(pos_util < pos_main);

  std::cout << "Composition logic verified." << std::endl;
}

void test_asset_composition() {
  std::cout << "Testing Asset-Based Shader Composition..." << std::endl;

  // Use test assets
  auto& sc = ShaderComposer::Get();

  size_t snippet_a_size;
  const char* snippet_a_code =
      (const char*)GetAsset(AssetId::ASSET_SHADER_SNIPPET_A, &snippet_a_size);
  assert(snippet_a_code != nullptr);
  sc.RegisterSnippet("SNIPPET_A", std::string(snippet_a_code, snippet_a_size));

  size_t snippet_b_size;
  const char* snippet_b_code =
      (const char*)GetAsset(AssetId::ASSET_SHADER_SNIPPET_B, &snippet_b_size);
  sc.RegisterSnippet("SNIPPET_B", std::string(snippet_b_code, snippet_b_size));

  std::string main_code =
      "fn main() -> f32 { return snippet_a() + snippet_b(); }";
  std::string result = sc.Compose({"SNIPPET_A", "SNIPPET_B"}, main_code);

  assert(result.find("fn snippet_a()") != std::string::npos);
  assert(result.find("fn snippet_b()") != std::string::npos);
  assert(result.find("fn main()") != std::string::npos);

  size_t pos_a = result.find("snippet_a");
  size_t pos_b = result.find("snippet_b");
  size_t pos_main = result.find("main");

  assert(pos_a < pos_b);
  assert(pos_b < pos_main);

  std::cout << "Asset-based composition logic verified." << std::endl;
}

int main() {
  test_composition();
  test_asset_composition();
  std::cout << "--- ALL SHADER COMPOSER TESTS PASSED ---" << std::endl;
  return 0;
}