// This file is part of the 64k demo project.
// It validates that all production shaders compile successfully with WebGPU.
// This catches issues like:
// - Invalid WGSL syntax (e.g., undefined functions like inverse())
// - Missing binding declarations
// - Type mismatches

#include "generated/assets.h"
#include "gpu/effects/shader_composer.h"
#include "gpu/effects/shaders.h"
#include "platform.h"
#include <cassert>
#include <cstdio>
#include <cstring>
#include <string>

static WGPUDevice g_device = nullptr;

// Initialize minimal WebGPU for shader compilation testing
static bool init_wgpu() {
  WGPUInstance instance = wgpuCreateInstance(nullptr);
  if (!instance) {
    fprintf(stderr, "Failed to create WGPU instance.\n");
    return false;
  }

  WGPURequestAdapterOptions adapter_opts = {};
  adapter_opts.powerPreference = WGPUPowerPreference_HighPerformance;

  WGPUAdapter adapter = nullptr;

#if defined(DEMO_CROSS_COMPILE_WIN32)
  auto on_adapter = [](WGPURequestAdapterStatus status, WGPUAdapter a,
                       const char* message, void* userdata) {
    if (status == WGPURequestAdapterStatus_Success) {
      *(WGPUAdapter*)userdata = a;
    }
  };
  wgpuInstanceRequestAdapter(instance, &adapter_opts, on_adapter, &adapter);
#else
  auto on_adapter = [](WGPURequestAdapterStatus status, WGPUAdapter a,
                       WGPUStringView message, void* userdata, void* user2) {
    (void)user2;
    (void)message;
    if (status == WGPURequestAdapterStatus_Success) {
      *(WGPUAdapter*)userdata = a;
    }
  };
  WGPURequestAdapterCallbackInfo adapter_cb = {};
  adapter_cb.mode = WGPUCallbackMode_WaitAnyOnly;
  adapter_cb.callback = on_adapter;
  adapter_cb.userdata1 = &adapter;
  wgpuInstanceRequestAdapter(instance, &adapter_opts, adapter_cb);
#endif

  // Try to wait for adapter (may not work on all platforms)
  for (int i = 0; i < 100 && !adapter; ++i) {
    wgpuInstanceProcessEvents(instance);
  }

  if (!adapter) {
    fprintf(stderr,
            "Warning: Could not get WGPU adapter (GPU compilation tests "
            "skipped)\n");
    return false;
  }

  WGPUDeviceDescriptor device_desc = {};

#if defined(DEMO_CROSS_COMPILE_WIN32)
  auto on_device = [](WGPURequestDeviceStatus status, WGPUDevice d,
                      const char* message, void* userdata) {
    if (status == WGPURequestDeviceStatus_Success) {
      *(WGPUDevice*)userdata = d;
    }
  };
  wgpuAdapterRequestDevice(adapter, &device_desc, on_device, &g_device);
#else
  auto on_device = [](WGPURequestDeviceStatus status, WGPUDevice d,
                      WGPUStringView message, void* userdata, void* user2) {
    (void)user2;
    (void)message;
    if (status == WGPURequestDeviceStatus_Success) {
      *(WGPUDevice*)userdata = d;
    }
  };
  WGPURequestDeviceCallbackInfo device_cb = {};
  device_cb.mode = WGPUCallbackMode_WaitAnyOnly;
  device_cb.callback = on_device;
  device_cb.userdata1 = &g_device;
  wgpuAdapterRequestDevice(adapter, &device_desc, device_cb);
#endif

  // Try to wait for device (may not work on all platforms)
  for (int i = 0; i < 100 && !g_device; ++i) {
    wgpuInstanceProcessEvents(instance);
  }

  if (!g_device) {
    fprintf(stderr,
            "Warning: Could not get WGPU device (GPU compilation tests "
            "skipped)\n");
    return false;
  }

  return true;
}

// Test shader compilation
static bool test_shader_compilation(const char* name,
                                     const char* shader_code) {
  printf("Testing compilation: %s...\n", name);

  if (!g_device) {
    printf("SKIPPED: %s (no GPU device)\n", name);
    return true;  // Not a failure, just skipped
  }

#if defined(DEMO_CROSS_COMPILE_WIN32)
  WGPUShaderModuleWGSLDescriptor wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderModuleWGSLDescriptor;
  wgsl_desc.code = shader_code;
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#else
  WGPUShaderSourceWGSL wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_desc.code = str_view(shader_code);
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#endif

  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(g_device, &shader_desc);

  if (!shader_module) {
    printf("FAILED: %s - shader compilation failed!\n", name);
    return false;
  }

  wgpuShaderModuleRelease(shader_module);
  printf("PASSED: %s\n", name);
  return true;
}

// Test composed shader with different modes
static bool test_composed_shader(const char* base_name, AssetId asset_id,
                                  bool with_bvh) {
  const char* mode_name = with_bvh ? "BVH" : "Linear";
  char test_name[128];
  snprintf(test_name, sizeof(test_name), "%s (%s mode)", base_name, mode_name);

  const char* shader_asset = (const char*)GetAsset(asset_id);
  std::string main_code = shader_asset;

  ShaderComposer::CompositionMap composition_map;
  if (with_bvh) {
    composition_map["render/scene_query_mode"] = "render/scene_query_bvh";
  } else {
    composition_map["render/scene_query_mode"] = "render/scene_query_linear";
  }

  std::string composed_shader =
      ShaderComposer::Get().Compose({}, main_code, composition_map);

  return test_shader_compilation(test_name, composed_shader.c_str());
}

int main() {
  printf("===========================================\n");
  printf("Shader Compilation Test Suite\n");
  printf("===========================================\n\n");

  bool gpu_available = init_wgpu();
  if (!gpu_available) {
    printf(
        "Note: GPU not available - running composition-only tests\n\n");
  }

  // Initialize shader composer
  InitShaderComposer();

  bool all_passed = true;

  // Test 1: Simple shaders that don't need composition
  printf("\n--- Test 1: Simple Shaders ---\n");
  all_passed &= test_shader_compilation(
      "Passthrough", (const char*)GetAsset(AssetId::ASSET_SHADER_PASSTHROUGH));
  all_passed &= test_shader_compilation(
      "Ellipse", (const char*)GetAsset(AssetId::ASSET_SHADER_ELLIPSE));
  all_passed &= test_shader_compilation(
      "Gaussian Blur",
      (const char*)GetAsset(AssetId::ASSET_SHADER_GAUSSIAN_BLUR));
  all_passed &= test_shader_compilation(
      "Solarize", (const char*)GetAsset(AssetId::ASSET_SHADER_SOLARIZE));

  // Test 2: Composed shaders (both BVH and Linear modes)
  printf("\n--- Test 2: Composed Shaders (BVH Mode) ---\n");
  all_passed &= test_composed_shader("Renderer 3D",
                                      AssetId::ASSET_SHADER_RENDERER_3D, true);
  all_passed &=
      test_composed_shader("Mesh Render", AssetId::ASSET_SHADER_MESH, true);

  printf("\n--- Test 3: Composed Shaders (Linear Mode) ---\n");
  all_passed &= test_composed_shader("Renderer 3D",
                                      AssetId::ASSET_SHADER_RENDERER_3D, false);
  all_passed &=
      test_composed_shader("Mesh Render", AssetId::ASSET_SHADER_MESH, false);

  // Test 3: Compute shaders
  printf("\n--- Test 4: Compute Shaders ---\n");
  all_passed &= test_shader_compilation(
      "Particle Compute",
      (const char*)GetAsset(AssetId::ASSET_SHADER_PARTICLE_COMPUTE));
  all_passed &= test_shader_compilation(
      "Particle Spray Compute",
      (const char*)GetAsset(AssetId::ASSET_SHADER_PARTICLE_SPRAY_COMPUTE));

  printf("\n===========================================\n");
  if (all_passed) {
    printf("All shader compilation tests PASSED ✓\n");
  } else {
    printf("Some shader compilation tests FAILED ✗\n");
  }
  printf("===========================================\n");

  if (g_device) {
    wgpuDeviceRelease(g_device);
  }

  return all_passed ? 0 : 1;
}