#include <glad/glad.h>
#include <GLFW/glfw3.h>
#include <vector>
#include <iostream>

// Vertex shader
const char* vertexShaderSource = R"(
#version 460 core
layout(location = 0) in vec2 position;

void main() {
    gl_Position = vec4(position, 0.0, 1.0);
}
)";

// Fragment shader using explicit sampler array
const char* fragmentShaderSource = R"(
#version 460 core

layout(location = 0) out vec4 fragColor;

// Array of sampler uniforms (bound to texture units 0-15)
layout(binding = 0) uniform sampler2DArray textures[16];

void main() {
    vec2 uv = gl_FragCoord.xy / vec2(800, 600);
    int texture_index = int(uv.x * 16);
    int array_element = int(uv.y * 4);
    fragColor = texture(textures[texture_index], vec3(uv, array_element));;
}
)";

GLuint createShaderProgram() {
  GLuint vertexShader = glCreateShader(GL_VERTEX_SHADER);
  glShaderSource(vertexShader, 1, &vertexShaderSource, nullptr);
  glCompileShader(vertexShader);

  GLuint fragmentShader = glCreateShader(GL_FRAGMENT_SHADER);
  glShaderSource(fragmentShader, 1, &fragmentShaderSource, nullptr);
  glCompileShader(fragmentShader);

  GLuint program = glCreateProgram();
  glAttachShader(program, vertexShader);
  glAttachShader(program, fragmentShader);
  glLinkProgram(program);

  GLint success;
  glGetProgramiv(program, GL_LINK_STATUS, &success);
  if (!success) {
    char infoLog[512];
    glGetProgramInfoLog(program, 512, nullptr, infoLog);
    std::cerr << "Shader linking failed:\n" << infoLog << std::endl;
  }

  glDeleteShader(vertexShader);
  glDeleteShader(fragmentShader);
  return program;
}

GLuint createTextureArray(unsigned char green) {
  GLuint texture;
  glCreateTextures(GL_TEXTURE_2D_ARRAY, 1, &texture);

  // Use DSA for all texture operations
  glTextureStorage3D(texture, 1, GL_SRGB8_ALPHA8, 256, 256, 4); // 4 layers

  // Generate simple procedural texture data
  std::vector<unsigned char> data(256 * 256 * 4 * 4); // RGBA8 * 4 layers
  for (int layer = 0; layer < 4; ++layer) {
    for (int y = 0; y < 256; ++y) {
      for (int x = 0; x < 256; ++x) {
        size_t idx = ((layer * 256 + y) * 256 + x) * 4;
        // Different color per layer
        data[idx + 0] = static_cast<unsigned char>(63 + layer * 64); // R
        data[idx + 1] = green;
        data[idx + 2] = static_cast<unsigned char>(63 + layer * 64); // B
        data[idx + 3] = 255;                                              // A
      }
    }
  }

  glTextureSubImage3D(texture, 0, 0, 0, 0, 256, 256, 4,
    GL_RGBA, GL_UNSIGNED_BYTE, data.data());

  // Modern texture parameters using DSA
  glTextureParameteri(texture, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
  glTextureParameteri(texture, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTextureParameteri(texture, GL_TEXTURE_WRAP_S, GL_REPEAT);
  glTextureParameteri(texture, GL_TEXTURE_WRAP_T, GL_REPEAT);
  glGenerateTextureMipmap(texture);

  return texture;
}

int main() {
  if (!glfwInit()) return -1;

  glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
  glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);
  glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
  glfwWindowHint(GLFW_OPENGL_DEBUG_CONTEXT, GLFW_TRUE);

  GLFWwindow* window = glfwCreateWindow(800, 600, "Modern OpenGL 4.6 Texture Arrays", nullptr, nullptr);
  if (!window) {
    glfwTerminate();
    return -1;
  }

  glfwMakeContextCurrent(window);
  gladLoadGLLoader((GLADloadproc)glfwGetProcAddress);

  // Create 16 texture arrays
  std::vector<GLuint> textureArrays(16);
  for (int i = 0; i < 16; ++i) {
    textureArrays[i] = createTextureArray(15 + i * 16);
  }

  // Create shader program
  GLuint program = createShaderProgram();

  // Get uniform location for the sampler array
  GLint texturesLoc = glGetUniformLocation(program, "textures");
  if (texturesLoc == -1) {
    std::cerr << "Uniform 'textures' not found in shader!" << std::endl;
    return -1;
  }

  // Create uniform array values (texture unit indices)
  GLint textureUnits[16];
  for (int i = 0; i < 16; ++i) {
    textureUnits[i] = i; // Texture unit 0-15
  }

  // Set the uniform array (only needs to be done once)
  glUseProgram(program);
  glUniform1iv(texturesLoc, 16, textureUnits);

  // Create simple fullscreen quad VAO with DSA
  GLuint vao, vbo;
  glCreateVertexArrays(1, &vao);
  glCreateBuffers(1, &vbo);

  float vertices[] = {
      -1.0f, -1.0f,
       1.0f, -1.0f,
      -1.0f,  1.0f,
       1.0f,  1.0f
  };

  glNamedBufferData(vbo, sizeof(vertices), vertices, GL_STATIC_DRAW);
  glVertexArrayVertexBuffer(vao, 0, vbo, 0, 2 * sizeof(float));
  glEnableVertexArrayAttrib(vao, 0);
  glVertexArrayAttribFormat(vao, 0, 2, GL_FLOAT, GL_FALSE, 0);
  glVertexArrayAttribBinding(vao, 0, 0);

  // Main render loop
  while (!glfwWindowShouldClose(window)) {
    glClear(GL_COLOR_BUFFER_BIT);

    glUseProgram(program);
    glBindVertexArray(vao);

    // Bind all 16 textures to consecutive texture units
    for (int i = 0; i < 16; ++i) {
      glBindTextureUnit(i, textureArrays[i]);
    }

    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

    glfwSwapBuffers(window);
    glfwPollEvents();
  }

  // Cleanup
  for (auto texture : textureArrays) {
    glDeleteTextures(1, &texture);
  }

  glDeleteBuffers(1, &vbo);
  glDeleteVertexArrays(1, &vao);
  glDeleteProgram(program);

  glfwDestroyWindow(window);
  glfwTerminate();
  return 0;
}