#include "renderer.h"
#include "gpu.h"
#include <iostream>
#define TINYOBJLOADER_IMPLEMENTATION
#include "tiny_obj_loader.h"


bool Renderer::init(SDL_Window* w) {
  win = w;
  dev = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, true, NULL);
  if (!dev) { std::cerr << "GPU device failed: " << SDL_GetError() << std::endl; return false; }

  SDL_ClaimWindowForGPUDevice(dev, win);

  load_obj();

  render_sampler = create_linear_sampler(dev);
  SDL_GPUSamplerCreateInfo depth_sampler_info {
    .min_filter    = SDL_GPU_FILTER_NEAREST,
    .mag_filter    = SDL_GPU_FILTER_NEAREST,
    .mipmap_mode   = SDL_GPU_SAMPLERMIPMAPMODE_NEAREST,
    .address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
    .address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
    .compare_op    = SDL_GPU_COMPAREOP_INVALID,
    .enable_compare = false,
  };
  depth_sampler = SDL_CreateGPUSampler(dev, &depth_sampler_info);
  
  
  resize_render_texture(960, 540);

  scene.create(dev, win);
  ray.create(dev, win);
  post.create(dev, win);

  return scene.pipeline && ray.pipeline;
}

void Renderer::destroy() {
  if (render_tex)     SDL_ReleaseGPUTexture(dev, render_tex);
  if (post_tex)       SDL_ReleaseGPUTexture(dev, post_tex);
  if (render_sampler) SDL_ReleaseGPUSampler(dev, render_sampler);
  if (vert_buff)      SDL_ReleaseGPUBuffer(dev, vert_buff);
  if (depth_tex)      SDL_ReleaseGPUTexture(dev, depth_tex);
  scene.destroy(dev);
  ray.destroy(dev);
  post.destroy(dev);
  SDL_DestroyGPUDevice(dev);
}

bool Renderer::load_obj() {
  const std::string& path = "assets/homer.obj";
  tinyobj::attrib_t attrib;
  std::vector<tinyobj::shape_t> shapes;
  std::vector<tinyobj::material_t> materials;
  std::string warn, err;
  if(!tinyobj::LoadObj(&attrib, &shapes, &materials, &warn, &err, path.c_str())) {
    std::cerr << "OBJ fail: " << warn << err << std::endl;
    return false;
  }

  vertices.clear();
  for(const auto& shape : shapes) {
    const auto& indices = shape.mesh.indices;
    for(size_t i = 0; i < indices.size(); i += 3) {
      // Use OBJ normals if present, otherwise compute face normal
      vec3 pos[3];
      vec3 nor[3];
      bool has_normals = true;

      for(int v = 0; v < 3; v++) {
        const auto& idx = indices[i + v];
        pos[v] = vec3(
          attrib.vertices[3 * idx.vertex_index + 0],
          attrib.vertices[3 * idx.vertex_index + 1],
          attrib.vertices[3 * idx.vertex_index + 2]
        );
        if(idx.normal_index >= 0) {
          nor[v] = vec3(
            attrib.normals[3 * idx.normal_index + 0],
            attrib.normals[3 * idx.normal_index + 1],
            attrib.normals[3 * idx.normal_index + 2]
          );
        } else {
          has_normals = false;
        }
      }

      if(!has_normals) {
        vec3 face_normal = glm::normalize(glm::cross(pos[1] - pos[0], pos[2] - pos[0]));
        nor[0] = nor[1] = nor[2] = face_normal;
      }

      for(int v = 0; v < 3; v++) {
        vertices.push_back(Vertex {
          .pos    = pos[v],
          .color  = vec4(abs(sin(pos[v].x)), abs(cos(pos[v].y)), abs(sin(pos[v].z) * cos(pos[v].y)), 1.0f),
          .normal = nor[v],
        });
      }
    }
  }

  SDL_GPUBufferCreateInfo info {
    .usage = SDL_GPU_BUFFERUSAGE_VERTEX,
    .size  = static_cast<uint32_t>(vertices.size() * sizeof(Vertex)),
  };
  vert_buff = SDL_CreateGPUBuffer(dev, &info);
  upload_buffer(dev, vert_buff, vertices.data(), info.size);
  return true;
}

void Renderer::resize_render_texture(uint32_t w, uint32_t h) {
  SDL_WaitForGPUIdle(dev);
  if (render_tex) SDL_ReleaseGPUTexture(dev, render_tex);
  if (post_tex) SDL_ReleaseGPUTexture(dev, post_tex);
  if (depth_tex) SDL_ReleaseGPUTexture(dev, depth_tex);

  SDL_GPUTextureCreateInfo depth_info {
    .type   = SDL_GPU_TEXTURETYPE_2D,
    .format = SDL_GPU_TEXTUREFORMAT_D32_FLOAT,
    .usage  = SDL_GPU_TEXTUREUSAGE_DEPTH_STENCIL_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER,
    .width  = w,
    .height = h,
    .layer_count_or_depth = 1,
    .num_levels = 1,
  };

  depth_tex = SDL_CreateGPUTexture(dev, &depth_info);
  render_tex = create_render_texture(dev, win, w, h);
  post_tex = create_render_texture(dev, win, w, h);
  render_w = w;
  render_h = h;
}


void Renderer::pass_scene(SDL_GPUCommandBuffer* cmd, const CameraUBO& ubo) {
  SDL_PushGPUVertexUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture     = render_tex,
    .clear_color = { 0.0, 0.0, 0.0, 1.0},
    .load_op     = SDL_GPU_LOADOP_CLEAR,
    .store_op    = SDL_GPU_STOREOP_STORE,
  };
  
  SDL_GPUDepthStencilTargetInfo depth_tgt {
    .texture          = depth_tex,
    .clear_depth      = 1.0f,
    .load_op          = SDL_GPU_LOADOP_CLEAR,
    .store_op         = SDL_GPU_STOREOP_STORE,
    .stencil_load_op  = SDL_GPU_LOADOP_DONT_CARE,
    .stencil_store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, &depth_tgt);
  SDL_BindGPUGraphicsPipeline(pass, scene.pipeline);
  SDL_GPUBufferBinding binding { .buffer = vert_buff, .offset = 0 };
  SDL_BindGPUVertexBuffers(pass, 0, &binding, 1);
  SDL_DrawGPUPrimitives(pass, vertices.size(), 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

void Renderer::pass_ray(SDL_GPUCommandBuffer* cmd, const CameraUBO& ubo) {
  SDL_PushGPUFragmentUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture  = post_tex,
    .load_op  = SDL_GPU_LOADOP_CLEAR,
    .store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, NULL);
  SDL_BindGPUGraphicsPipeline(pass, ray.pipeline);
  SDL_GPUTextureSamplerBinding texs[2] = {
    { .texture = render_tex, .sampler = render_sampler },
    { .texture = depth_tex, .sampler = depth_sampler },
  };
  SDL_BindGPUFragmentSamplers(pass, 0, texs, 2);
  SDL_DrawGPUPrimitives(pass, 3, 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

void Renderer::pass_post(SDL_GPUCommandBuffer* cmd, SDL_GPUTexture* swapchain, const CameraUBO& ubo) {
  SDL_PushGPUFragmentUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture  = swapchain,
    .load_op  = SDL_GPU_LOADOP_CLEAR,
    .store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, NULL);
  SDL_BindGPUGraphicsPipeline(pass, post.pipeline);
  SDL_GPUTextureSamplerBinding tex { .texture = post_tex, .sampler = render_sampler };
  SDL_BindGPUFragmentSamplers(pass, 0, &tex, 1);
  SDL_DrawGPUPrimitives(pass, 3, 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

static int neighbors(const std::array<uint8_t, 256*256>& grid, int x, int y) {
  int count = 0;
  for (int dy = -1; dy <= 1; dy++)
    for (int dx = -1; dx <= 1; dx++)
      if (dx || dy)
        count += grid[((y + dy + 256) & 255) * 256 + ((x + dx + 256) & 255)];
  return count;
}

std::array<std::array<uint8_t, 256*256>, 2> graph;

static void gol(std::array<std::array<uint8_t, 256*256>, 2>& graph, size_t frame) {
  size_t src = (frame >> 8) & 1;
  size_t dst = src ^ 1;
  for (int y = 0; y < 256; y++) {
    for (int x = 0; x < 256; x++) {
      int n = neighbors(graph[src], x, y);
      bool alive = graph[src][y * 256 + x];
      graph[dst][y * 256 + x] = alive ? (n == 2 || n == 3) : (n == 3);
    }
  }
}

static vec4 int_to_color(int i) {
  uint32_t x = static_cast<uint32_t>(i);

  x ^= x >> 17;
  x *= 0xed5ad4bb;
  x ^= x >> 11;
  x *= 0xac4c1b51;
  x ^= x >> 15;
  x *= 0x31848bab;
  x ^= x >> 14;

  float r = ((x >>  0) & 0xFF) / 255.0f;
  float g = ((x >>  8) & 0xFF) / 255.0f;
  float b = ((x >> 16) & 0xFF) / 255.0f;

  return vec4(r, g, b, 1.0f);
}

void Renderer::add_gol_layer() {
  size_t layerno = frame >> 8;

  if (frame == 0) {
    for (size_t i = 0; i < 256 * 256; i++) {
      graph[0][i] = rand() & 1;
    }
  }


  gol(graph, frame);

  size_t src = ((frame >> 8) & 1) ^ 1;
  float height = layerno * 10.0f;

  for (int y = 0; y < 256; y++) {
    for (int x = 0; x < 256; x++) {
      if (graph[src][y * 256 + x]) {
        float s = 0.01f;

        vec3 p = vec3(x + s * sinf((float)frame), height, y + cosf((float)frame)) * s * 3.0f;

        vec3 v0 = p + vec3(-s, -s, 0);
        vec3 v1 = p + vec3( s, -s, 0);
        vec3 v2 = p + vec3( s,  s, 0);
        vec3 v3 = p + vec3(-s,  s, 0);
        const vec4 color = int_to_color(layerno);
        const vec3 normal = vec3(0.0, 1.0, 0.0);


        vertices.push_back(Vertex {v0, color, normal});
        vertices.push_back(Vertex {v1, color, normal});
        vertices.push_back(Vertex {v2, color, normal});

        vertices.push_back(Vertex {v2, color, normal});
        vertices.push_back(Vertex {v3, color, normal});
        vertices.push_back(Vertex {v0, color, normal});
      }
    }
  }
  if (vert_buff) SDL_ReleaseGPUBuffer(dev, vert_buff);

  SDL_GPUBufferCreateInfo info {
    .usage = SDL_GPU_BUFFERUSAGE_VERTEX,
    .size  = static_cast<uint32_t>(vertices.size() * sizeof(Vertex)),
  };

  vert_buff = SDL_CreateGPUBuffer(dev, &info);
  upload_buffer(dev, vert_buff, vertices.data(), info.size);
}


void Renderer::draw(const CameraUBO& ubo) {
  if(frame % 256 == 0 && (frame >> 8) < 20) {
    add_gol_layer();
  }

  SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(dev);

  SDL_GPUTexture* swapchain;
  uint32_t sw_w, sw_h;
  SDL_WaitAndAcquireGPUSwapchainTexture(cmd, win, &swapchain, &sw_w, &sw_h);

  if (!swapchain) {
    SDL_SubmitGPUCommandBuffer(cmd);
    return;
  }

  if (sw_w != render_w || sw_h != render_h)
    resize_render_texture(sw_w, sw_h);

  pass_scene(cmd, ubo);
  pass_ray(cmd, ubo);
  pass_post(cmd, swapchain, ubo);

  SDL_SubmitGPUCommandBuffer(cmd);
  frame++;
}