#define GLEW_STATIC
#include <GL/glew.h>
#include <GLFW/glfw3.h>
#include <SOIL/SOIL.h>
#include <iostream>
#include <vector>
#include <malloc.h>
#include <chrono>

unsigned char* _in;
unsigned char* _out;
GLuint texture_in;
GLuint texture_out;

void loadTexture(int& width, int& height) {
  width = 1920 << 1;
  height = 1080 << 1;

  _in = (unsigned char*)malloc(width * height * 3);
  _out = (unsigned char*)malloc(width * height * 3);
//  for (int y = 0; y < height; ++y) {
//    for (int x = 0; x < width; ++x) {
//      ((int*)image)[y * width + x] = (((x ^ y) & 0xff) << 24 | (x & 0xff << 16) | (y & 0xff << 8) | ((x ^ y) & 0xff)) ^ 0x7f7f7f7f;
//    }
//  }

  uint8_t* p = _in;
  for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
      if (x < 16) {
        *p++ = 0x80;
        *p++ = 0xff;
        *p++ = 0xff;
      } else if (x >= 64 && x < 192 && y >= 64 && y < 192) {
        if (x < 80 || y < 80) {
          *p++ = 0xb5;
          *p++ = 0xb5;
          *p++ = 0xb5;
        } else if ((x ^ y) & 16) {
          *p++ = 0;
          *p++ = 0xff;
          *p++ = 0;
        } else {
          *p++ = 0xff;
          *p++ = 0;
          *p++ = 0xff;
        }
      } else {
        int m = 1;
        if (y >= 128) m <<= 1;
        if (x >= 128) m <<= 2;
        if ((x ^ y) & m) {
          *p++ = 0xff;
          *p++ = 0xff;
          *p++ = 0xff;
        } else {
          *p++ = 0x00;
          *p++ = 0x00;
          *p++ = 0x00;
        }
      }
    }
  }

  glGenTextures(1, &texture_in);
  glGenTextures(1, &texture_out);

  glBindTexture(GL_TEXTURE_2D, texture_in);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, _in);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);

  glBindTexture(GL_TEXTURE_2D, texture_out);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, nullptr);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);

//  free(image);
//  SOIL_free_image_data(image);
  glBindTexture(GL_TEXTURE_2D, 0);
}




GLuint createComputeShader(const char* shaderSource) {
  GLuint shader = glCreateShader(GL_COMPUTE_SHADER);
  glShaderSource(shader, 1, &shaderSource, nullptr);
  glCompileShader(shader);

  GLint success;
  glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
  if (!success) {
    char infoLog[512];
    glGetShaderInfoLog(shader, 512, nullptr, infoLog);
    std::cerr << "ERROR::SHADER::COMPUTE::COMPILATION_FAILED\n" << infoLog << std::endl;
  }

  GLuint program = glCreateProgram();
  glAttachShader(program, shader);
  glLinkProgram(program);

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

  glDeleteShader(shader);
  return program;
}











#include <iostream>
#include <fstream>
#include <sstream>

void readTextureToCPU(GLuint texture, int width, int height, bool save = false) {
  glBindTexture(GL_TEXTURE_2D, texture);

//  std::vector<unsigned char> data(width * height * 3);
  glGetTexImage(GL_TEXTURE_2D, 0, GL_RGB, GL_UNSIGNED_BYTE, _out);

  if (save) {
    std::ofstream out("out.pnm", std::ios::binary);
    out << "P6 " << width << " " << height << " 255\n";
    out.write((char*)_out, width * height * 3);
  }

  // Do something with the data
 // std::cout << "First pixel: ("
 //   << static_cast<int>(data[0]) << ", "
 //   << static_cast<int>(data[1]) << ", "
 //   << static_cast<int>(data[2]) << ", "
 //   << static_cast<int>(data[3]) << ")" << std::endl;

  glBindTexture(GL_TEXTURE_2D, 0);
}






int main() {
    // Initialize GLFW and create a window
  if (!glfwInit()) {
    std::cerr << "Failed to initialize GLFW" << std::endl;
    return -1;
  }

  glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE);
  GLFWwindow* window = glfwCreateWindow(800, 600, "OpenGL Compute Shader Example", nullptr, nullptr);
  if (!window) {
    std::cerr << "Failed to create GLFW window" << std::endl;
    glfwTerminate();
    return -1;
  }

  glfwMakeContextCurrent(window);

  // Initialize GLEW
  if (glewInit() != GLEW_OK) {
    std::cerr << "Failed to initialize GLEW" << std::endl;
    return -1;
  }

  // Load the texture
  int width, height;
  loadTexture(width, height);
  if (!texture_in || !texture_out) {
    throw("Texture problem");
    return -1;
  }

  // Create and use the compute shader
  std::ifstream in("test.glsl");
  if (!in.is_open()) throw("File did not open");
  std::stringstream ss;
  ss << in.rdbuf();
  GLuint computeProgram = createComputeShader(ss.str().c_str());

//  glBindTexture(GL_TEXTURE_2D, texture_in);
  glUseProgram(computeProgram);


  glBindImageTexture(0, texture_in, 0, GL_FALSE, 0, GL_READ_ONLY, GL_RGBA8);
  glBindImageTexture(1, texture_out, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA8);

  // Dispatch the compute shader
  glDispatchCompute((GLuint)ceil(width / 32.0), (GLuint)ceil(height / 32.0), 1);

// Wait for the compute shader to finish
  glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // Read the scrambled texture back to the CPU

  readTextureToCPU(texture_out, width, height, true);

  glUseProgram(0);


  // Clean up
  glDeleteProgram(computeProgram);
  glDeleteTextures(1, &texture_in);
  glDeleteTextures(1, &texture_out);
  glfwDestroyWindow(window);
  glfwTerminate();

  return 0;
}