#include "aliases.h"
#include "img.h"
#include "obj.h"
#include "render.h"
#include "shader.h"
#include "utils.h"
#include "vec.h"
#include <math.h>

typedef struct shader Shader;
struct shader {
  M4x4f model_view;
  M4x4f projection;
};

ShaderID perspective_phong = {0};
ShaderID perspective_albedo = {0};
ShaderID orthographic_phong = {0};
ShaderID orthographic_albedo = {0};

internal Shader perspective = {0};
internal Shader orthographic = {0};

internal V3f g_ambient_light = {0.6f, 0.45f, 0.55f};
internal PointLight g_directional_light = {
    .diffuse_intensity = {0.9f, 1.0f, 1.0f},
    .specular_intensity = {1.0f, 0.8f, 2.0f},
    .position = {1.05f, 0.9f, 1.2f},
};
internal V3f g_eye = {0.2f, 0.1f, 0.75f};
internal V3f g_target = {0};
internal V3f g_up = {0.0f, 1.0f, 0.0f};

internal V3f general_shader_vertex(void *shader, const V3f *vertex);
internal FragmentResult phong_shader_fragment(void *shader,
                                              const FragmentData *data,
                                              const Colour *colour,
                                              const Model *model);
internal FragmentResult albedo_shader_fragment(void *shader,
                                               const FragmentData *data,
                                               const Colour *colour,
                                               const Model *model);
internal M4x4f get_projection_matrix(ProjectionType projection_type);
internal f32 get_intensity(const V3f *normal);

void load_shaders(void) {
  M4x4f model_view = lookat(g_eye, g_target, g_up);
  M4x4f orthographic_projection =
      get_projection_matrix(PROJECTION_TYPE_ORTHOGRAPHIC);
  M4x4f perspective_projection =
      get_projection_matrix(PROJECTION_TYPE_PERSPECTIVE);

  perspective.model_view = orthographic.model_view = model_view;
  perspective.projection = perspective_projection;
  orthographic.projection = orthographic_projection;

  perspective_phong = register_shader(&perspective, general_shader_vertex,
                                      phong_shader_fragment);
  perspective_albedo = register_shader(&perspective, general_shader_vertex,
                                       albedo_shader_fragment);
  orthographic_phong = register_shader(&orthographic, general_shader_vertex,
                                       phong_shader_fragment);
  orthographic_albedo = register_shader(&orthographic, general_shader_vertex,
                                        albedo_shader_fragment);
}

internal V3f general_shader_vertex(void *shader, const V3f *vertex) {
  Shader *shader_ptr = (Shader *)shader;

  V4f vh = {.x = vertex->x, .y = vertex->y, .z = vertex->z, .w = 1.0f};
  vh = mat4x4_mul_vec4(shader_ptr->projection,
                       mat4x4_mul_vec4(shader_ptr->model_view, vh));

  return project_vec4(vh);
}

internal FragmentResult phong_shader_fragment(void *shader,
                                              const FragmentData *data,
                                              const Colour *colour,
                                              const Model *model) {
  Shader *shader_ptr = (Shader *)shader;

  V4f hnorm;
  V3f norm;
  if (model->normal) {
    u64 nm_x = data->tex_coords.u * model->normal->width;
    u64 nm_y = data->tex_coords.v * model->normal->height;

    Colour pixel = get_pixel(Colour, model->normal, nm_x, nm_y);
    hnorm = (V4f){
        .x = pixel.r,
        .y = pixel.g,
        .z = pixel.b,
        .w = 1.0f,
    };
  } else {
    hnorm = V3_to_V4(data->normal);
  }

  M4x4f matrix = mat4x4_mul(shader_ptr->projection, shader_ptr->model_view);
  M4x4f transposed = mat4x4_transpose(matrix);
  M4x4f inv_transpose = mat4x4_inv(transposed);

  hnorm = mat4x4_mul_vec4(inv_transpose, hnorm);
  norm = project_vec4(hnorm);
  normalise_v3(norm);

  Colour output;
  if (model->texture) {
    u64 tx_x = data->tex_coords.u * model->texture->width;
    u64 tx_y = data->tex_coords.v * model->texture->height;
    output = get_pixel(Colour, model->texture, tx_x, tx_y);
  } else {
    output = *colour;
  }

  V3f local_colour = {.r = output.r, .g = output.g, .b = output.b};
  local_colour = num_div_v3(local_colour, 255.0f);

  // Ambient term
  V3f intensity = num_mul_v3(g_ambient_light, model->material.ambient);

  V4f hdir = V3_to_V4(g_directional_light.position);
  hdir = mat4x4_mul_vec4(matrix, hdir);
  V3f light_pos = project_vec4(hdir);
  normalise_v3(light_pos);
  V3f light_dir = sub_v3(light_pos, data->position);

  // Diffuse term
  f32 l_dot_n = dot_v3(norm, light_dir);
  if (l_dot_n <= 0.0f) {
    goto RETURN_OUTPUT_COLOUR;
  }

  V3f diffuse = num_mul_v3(g_directional_light.diffuse_intensity,
                           model->material.diffuse * l_dot_n);
  intensity = add_v3(intensity, diffuse);

  // Specular term
  V3f _2_l_dot_n_norm = num_mul_v3(norm, 2.0f * l_dot_n);
  V3f reflected = sub_v3(_2_l_dot_n_norm, light_dir);
  normalise_v3(reflected);
  V3f v = sub_v3(g_eye, data->position);
  normalise_v3(v);
  f32 r_dot_v = dot_v3(reflected, v);
  if (r_dot_v <= 0.0f) {
    goto RETURN_OUTPUT_COLOUR;
  }

  V3f specular = num_mul_v3(g_directional_light.specular_intensity,
                            model->material.specular *
                                powf(r_dot_v, model->material.shininess));
  intensity = add_v3(intensity, specular);

RETURN_OUTPUT_COLOUR:
  intensity.r = clamp(intensity.r, 0.0f, 1.0f);
  intensity.g = clamp(intensity.g, 0.0f, 1.0f);
  intensity.b = clamp(intensity.b, 0.0f, 1.0f);

  local_colour = mul_v3(local_colour, intensity);

  output.r = 255.0f * local_colour.r;
  output.g = 255.0f * local_colour.g;
  output.b = 255.0f * local_colour.b;

  return (FragmentResult){.colour = output};
}

internal FragmentResult albedo_shader_fragment(void *shader,
                                               const FragmentData *data,
                                               const Colour *colour,
                                               const Model *model) {
  return (FragmentResult){.colour = *colour};
}

internal M4x4f get_projection_matrix(ProjectionType projection_type) {
  if (projection_type == PROJECTION_TYPE_PERSPECTIVE) {
    // Calculate projection matrix
    V3f cam = V3(V3f, f32, g_target.x, g_target.y, g_target.z, g_eye.x, g_eye.y,
                 g_eye.z);
    normalise_v3(cam);
    f32 coeff = -1.0f / magnitude_v3(cam) * 0.5f;
    return projection(coeff);
  }

  return mat4x4_identity;
}

internal f32 get_intensity(const V3f *normal) {
  V3f light = g_directional_light.position;
  normalise_v3(light);

  f32 intensity = dot_v3((*normal), light);
  if (intensity < 0.0f) {
    intensity = 0.001f;
  }

  return intensity;
}