#include <renderkid/rasterization.hpp>
#include <renderkid/vertex_stage.hpp>

#define GLM_ENABLE_EXPERIMENTAL

#include <glm/gtx/string_cast.hpp>
#include <glm/ext/scalar_common.hpp>
#include <glm/vec4.hpp>
#include <raylib.h>
#include <utility>
#include <cstdint>
#include <memory>
#include <vector>
#include <cmath>
#include <map>

std::allocator<float> depth_buffer_allocator;
std::allocator<uint8_t> texture_allocator;


inline float determinant(const glm::vec4 &a, const glm::vec4 &b, const glm::vec4 &c) {
  glm::vec4 ab = glm::vec4(b);
  glm::vec4 ac = glm::vec4(c);

  ab -= a;
  ac -= a;

  return ab.y * ac.x - ab.x * ac.y;
}


_raster::Rasterizer::Rasterizer(int width, int height) {
  this->width = width;
  this->height = height;
  textures = std::map<uint8_t, _raster::UserTexture>();

  int depth_buffer_size = width * height;

  depth_buffer = depth_buffer_allocator.allocate(depth_buffer_size);

  for (int i = 0; i < depth_buffer_size; ++i) {
    depth_buffer[i] = INFINITY;
  }

  next_texture_id = 0;
  framebuffer = GenImageColor(width, height, BLACK);
  framebuffer_texture = LoadTextureFromImage(framebuffer);
  SetTextureFilter(framebuffer_texture, TEXTURE_FILTER_POINT);
}


uint8_t _raster::Rasterizer::create_texture(int width, int height, int n_channels) {
  _raster::UserTexture new_texture;
  new_texture.width = width;
  new_texture.height = height;
  new_texture.channels = n_channels;
  new_texture.data = texture_allocator.allocate(width * height * n_channels);

  textures[next_texture_id] = new_texture;

  return next_texture_id++;
}


void _raster::Rasterizer::destroy_texture(uint8_t id) {
  int texture_size = textures[id].width * textures[id].height * textures[id].channels;
  texture_allocator.deallocate(textures[id].data, texture_size);
  textures.erase(id);
}


void _raster::Rasterizer::load_data_to_texture(uint8_t *data) {
  _raster::UserTexture texture = textures[current_texture];
  int texture_size = texture.width * texture.height * texture.channels;

  memcpy(textures[current_texture].data, data, texture_size);
}


inline bool is_left_or_top_edge(const glm::vec4 &start, const glm::vec4 &end) {
  glm::vec4 edge = end;
  edge -= start;

  return (edge.y > 0) || (edge.y == 0 && edge.x < 0);
}


glm::vec4 _gen_frag_color_from_vertex_color(const glm::vec4 &frag_coord,
                                            const glm::vec4 &a_color,
                                            const glm::vec4 &b_color,
                                            const glm::vec4 &c_color) {
  glm::vec4 color = a_color * frag_coord.x + b_color * frag_coord.y + c_color * frag_coord.z;

  float color_a = glm::max(0, glm::min(255, (int)glm::floor(color.x * 256.0f)));
  float color_b = glm::max(0, glm::min(255, (int)glm::floor(color.y * 256.0f)));
  float color_c = glm::max(0, glm::min(255, (int)glm::floor(color.z * 256.0f)));

  return glm::vec4(color_a, color_b, color_c, 255.0);
}


glm::vec4 _raster::Rasterizer::gen_frag_color_from_texture(const glm::vec4 &frag_coord,
                                                           const glm::vec4 &a_coord,
                                                           const glm::vec4 &b_coord,
                                                           const glm::vec4 &c_coord) {
  glm::vec4 interpolated_uv = frag_coord.x * a_coord + frag_coord.y * b_coord + frag_coord.z * c_coord;
  float interpolated_w = frag_coord.x * a_coord.w + frag_coord.y * b_coord.w + frag_coord.z * c_coord.w;

  _raster::UserTexture texture = textures[current_texture];
  int texture_size = texture.width * texture.height * texture.channels;

  int coord_x = (int)((interpolated_uv.x / interpolated_w) * texture.width) % texture.width;
  int coord_y = (int)((interpolated_uv.y / interpolated_w) * texture.height) % texture.height;
  int idx = (coord_y * texture.width + coord_x) * texture.channels;

  return glm::vec4(texture.data[idx], texture.data[idx + 1], texture.data[idx + 2], texture.data[idx + 3]);
}


glm::vec4 _raster::Rasterizer::gen_frag_color_from_attributes(const glm::vec4 &frag_coord, const _vertex::Vertex &a,
                                                              const _vertex::Vertex &b, const _vertex::Vertex &c) {
  glm::vec4 result = glm::vec4(1.0f);

  for (size_t i = 0; i < MAX_ATTRIBS; ++i) {
    switch (a.attribs[i].type) {
      case 0:
        return result;
      case COLOR_RGB_ATTRIBUTE:
      case COLOR_RGBA_ATTRIBUTE:
        result *= _gen_frag_color_from_vertex_color(frag_coord, a.attribs[i].data,
                                                    b.attribs[i].data, c.attribs[i].data);
        break;
      case TEXTURE_COORD_ATTRIBUTE:
        result *= gen_frag_color_from_texture(frag_coord, a.attribs[i].data,
                                              b.attribs[i].data, c.attribs[i].data);
        break;
    }
  }

  return result;
}


void _raster::Rasterizer::raster_triangle(const _vertex::Primitive &triangle) {
  glm::vec4 a = triangle.vertices[0].coords;
  glm::vec4 b = triangle.vertices[1].coords;
  glm::vec4 c = triangle.vertices[2].coords;

  const float xmin = glm::fmin(a.x, b.x, c.x);
  const float ymin = glm::fmin(a.y, b.y, c.y);
  const float xmax = glm::fmax(a.x, b.x, c.x);
  const float ymax = glm::fmax(a.y, b.y, c.y);

  for (float y = ymin; y <= ymax; ++y) {
    if (y < 0 || y > height) {
      continue;
    }

    for (float x = xmin; x <= xmax; ++x) {
      if (x < 0 || x > width) {
        continue;
      }

      glm::vec4 pixel = glm::vec4(x, y, 0.0f, 0.0f);

      float area = determinant(a, b, c);
      bool clockwise = false;

      if (area < 0) {
        std::swap(b, c);
        area = -area;
        clockwise = true;
      }

      float w0 = determinant(b, c, pixel);
      float w1 = determinant(c, a, pixel);
      float w2 = determinant(a, b, pixel);

      if (is_left_or_top_edge(b, c)) {
        --w0;
      }

      if (is_left_or_top_edge(c, a)) {
        --w1;
      }

      if (is_left_or_top_edge(a, b)) {
        --w2;
      }

      if (w0 >= 0 && w1 >= 0 && w2 >= 0) {
        glm::vec4 bary_coords = glm::vec4((w0 + 1) / area, (w1 + 1) / area, (w2 + 1) / area, 0.0f);

        float frag_z = 1 / (a.z * bary_coords.x + b.z * bary_coords.y + c.z * bary_coords.z);
        glm::vec4 frag_color;

        frag_color = gen_frag_color_from_attributes(
          bary_coords,
          triangle.vertices[0],
          triangle.vertices[1],
          triangle.vertices[2]
        );

        int buffer_index = ((int)y) * width + ((int)x);
        float previous_z = depth_buffer[buffer_index];

        if (frag_z < previous_z) {
          depth_buffer[buffer_index] = frag_z;

          ImageDrawPixel(&framebuffer, x, y, (Color){
            static_cast<uint8_t>(frag_color.x),
            static_cast<uint8_t>(frag_color.y),
            static_cast<uint8_t>(frag_color.z),
            static_cast<uint8_t>(frag_color.w)
          });
        }
      }
    }
  }
}


void _raster::Rasterizer::draw_framebuffer() {
  UpdateTexture(framebuffer_texture, framebuffer.data);
  DrawTexture(framebuffer_texture, 0, 0, WHITE);
}


void _raster::Rasterizer::clear_framebuffer() {
  ImageClearBackground(&framebuffer, BLACK);
}


void _raster::Rasterizer::draw_primitives(uint8_t primitive_type, std::vector<_vertex::Primitive> &primitives) {
  BeginDrawing();
  clear_framebuffer();
  int depth_buffer_size = width * height;

  for (int i = 0; i < depth_buffer_size; ++i) {
    depth_buffer[i] = INFINITY;
  }

  for (const _vertex::Primitive &primitive : primitives) {
    raster_triangle(primitive);
  }

  primitives.clear();
  draw_framebuffer();
  EndDrawing();
}