#include <stdio.h>
#include <stdbool.h>
#include "shader.h"
#include "mons_math/vec4.h"
#include "mons_math/vec3.h"
#include "mons_math/vec2.h"
#include "mons_math/mat4.h"
#include "hashmap.h"
#include <stdlib.h>

typedef enum mons_uniform_type {
    MONS_UNIFORM_INT,
    MONS_UNIFORM_FLOAT,
    MONS_UNIFORM_VEC2,
    MONS_UNIFORM_VEC3,
    MONS_UNIFORM_VEC4,
    MONS_UNIFORM_MAT4,
} mons_uniform_type;

typedef struct mons_uniform_m4 {
    mons_mat4 matrix;
    bool transpose;
} mons_uniform_m4;

typedef union mons_uniform_data {
    int i;
    float f;
    mons_vec2 v2;
    mons_vec3 v3;
    mons_vec4 v4;
    mons_uniform_m4 m4;
} mons_uniform_data;

typedef struct mons_uniform {
    mons_uniform_type type;
    mons_uniform_data data;
} mons_uniform;

static mons_hashmap *GLOBAL_UNIFORMS = NULL;
mons_hashmap *get_global_uniforms() {
    if (GLOBAL_UNIFORMS == NULL) {
        GLOBAL_UNIFORMS = malloc(sizeof(mons_hashmap));
        *GLOBAL_UNIFORMS = mons_hashmap_new(sizeof(mons_uniform), 10);
    }
    return GLOBAL_UNIFORMS;
}

void mons_shader_apply_global_uniforms(mons_program shader) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    for (int i = 0; i < global_uniforms->len; i++) {
        mons_hashmap_pair pair;
        mons_hashmap_at(*global_uniforms, i, &pair);
        char *name = pair.key;
        mons_uniform uniform = *(mons_uniform *)pair.value;
        switch (uniform.type) {
        case MONS_UNIFORM_INT:
            // TODO:
            break;
        case MONS_UNIFORM_FLOAT:
            mons_shader_set_float(shader, name, uniform.data.f);
            break;
        case MONS_UNIFORM_VEC2:
            mons_shader_set_vec2(shader, name, uniform.data.v2);
            break;
        case MONS_UNIFORM_VEC3:
            mons_shader_set_vec3(shader, name, uniform.data.v3);
            break;
        case MONS_UNIFORM_VEC4:
            mons_shader_set_vec4(shader, name, uniform.data.v4);
            break;
        case MONS_UNIFORM_MAT4:
            mons_shader_set_mat4(shader, name, uniform.data.m4.matrix, uniform.data.m4.transpose);
            break;
        }
    }
}

void mons_shader_set_float_global(char *uniform_name, float value) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    mons_uniform uniform = {
        .type = MONS_UNIFORM_FLOAT,
        .data.f = value,
    };
    mons_hashmap_insert(global_uniforms, uniform_name, &uniform);
}

void mons_shader_set_vec2_global(char *uniform_name, mons_vec2 value) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    mons_uniform uniform = {
        .type = MONS_UNIFORM_VEC2,
        .data.v2 = value,
    };
    mons_hashmap_insert(global_uniforms, uniform_name, &uniform);
}

void mons_shader_set_vec3_global(char *uniform_name, mons_vec3 value) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    mons_uniform uniform = {
        .type = MONS_UNIFORM_VEC3,
        .data.v3 = value,
    };
    mons_hashmap_insert(global_uniforms, uniform_name, &uniform);
}

void mons_shader_set_vec4_global(char *uniform_name, mons_vec4 value) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    mons_uniform uniform = {
        .type = MONS_UNIFORM_VEC4,
        .data.v4 = value,
    };
    mons_hashmap_insert(global_uniforms, uniform_name, &uniform);
}

void mons_shader_set_mat4_global(char *uniform_name, mons_mat4 value, bool transpose) {
    mons_hashmap *global_uniforms = get_global_uniforms();
    mons_uniform uniform = {
        .type = MONS_UNIFORM_MAT4,
        .data.m4 = {
            .matrix = value,
            .transpose = transpose,
        },
    };
    mons_hashmap_insert(global_uniforms, uniform_name, &uniform);
}

mons_shader mons_create_shader(mons_shader_type shader_type, char *source,
                               int source_length) {
    unsigned int shader = glCreateShader(shader_type);
    const char *shader_source = (const char *)source;
    int shader_len = (int)source_length;
    glShaderSource(shader, 1, &shader_source, &shader_len);
    glCompileShader(shader);
    int success;
    char info_log[512];
    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        glGetShaderInfoLog(shader, 512, NULL, info_log);
        printf("Shader Compilation Failed: %s\n", info_log);
    }
    return shader;
}

mons_program mons_create_program(mons_shader vertex_shader,
                                 mons_shader fragment_shader) {
    unsigned int shader_program = glCreateProgram();
    glAttachShader(shader_program, vertex_shader);
    glAttachShader(shader_program, fragment_shader);
    glLinkProgram(shader_program);
    int success;
    char info_log[512];
    glGetProgramiv(shader_program, GL_LINK_STATUS, &success);
    if (!success) {
        glGetProgramInfoLog(shader_program, 512, NULL, info_log);
    }
    return shader_program;
}

void mons_shader_set_float(mons_program shader, char *uniform_name,
                           float value) {
    glUniform1f(glGetUniformLocation(shader, uniform_name), value);
}

void mons_shader_set_vec2(mons_program shader, char *uniform_name,
                          mons_vec2 value) {
    glUniform2fv(glGetUniformLocation(shader, uniform_name), 1,
                 (float *)&value);
}

void mons_shader_set_vec3(mons_program shader, char *uniform_name,
                          mons_vec3 value) {
    glUniform3fv(glGetUniformLocation(shader, uniform_name), 1,
                 (float *)&value);
}

void mons_shader_set_vec4(mons_program shader, char *uniform_name,
                          mons_vec4 value) {
    glUniform4fv(glGetUniformLocation(shader, uniform_name), 1,
                 (float *)&value);
}

void mons_shader_set_mat4(mons_program shader, char *uniform_name,
                          mons_mat4 value, bool transpose) {
    glUniformMatrix4fv(glGetUniformLocation(shader, uniform_name), 1, transpose,
                       (float *)&value);
}