/*
 * Tuxánci 2 - A first person shooter
 * Copyright (C) 2025-2026  Connor Thomson
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include "model.h"
#include "camera.h"
#include "files.h"
#include "ta_math.h"

static bool ta_model_push_vertex(ta_model_mesh *mesh, size_t *capacity, const ta_model_vertex *vertex) {
    if (mesh->vertex_count == *capacity) {
        size_t next_capacity = *capacity == 0 ? 1024 : *capacity * 2;
        ta_model_vertex *vertices = realloc(mesh->vertices, next_capacity * sizeof(*vertices));
        if (!vertices) {
            return false;
        }
        mesh->vertices = vertices;
        *capacity = next_capacity;
    }
    mesh->vertices[mesh->vertex_count++] = *vertex;
    return true;
}

bool ta_model_build(const float positions[][TA_MODEL_DIMENSION], size_t position_count, const unsigned int faces[][TA_MODEL_DIMENSION], size_t face_count, ta_model_mesh *mesh) {
    float *normal_accum = NULL;
    size_t vertex_capacity = 0;
    bool success = false;

    memset(mesh, 0, sizeof(*mesh));
    if (position_count == 0 || face_count == 0) {
        return false;
    }

    for (size_t index = 0; index < face_count; index++) {
        for (int corner = 0; corner < TA_MODEL_DIMENSION; corner++) {
            if (faces[index][corner] >= position_count) {
                return false;
            }
        }
    }

    normal_accum = calloc(position_count * TA_MODEL_DIMENSION, sizeof(*normal_accum));
    if (!normal_accum) {
        goto done;
    }

    for (size_t index = 0; index < face_count; index++) {
        const float *first  = positions[faces[index][0]];
        const float *second = positions[faces[index][1]];
        const float *third  = positions[faces[index][2]];
        float edge_a[TA_MODEL_DIMENSION] = {
            second[0] - first[0],
            second[1] - first[1],
            second[2] - first[2]
        };
        float edge_b[TA_MODEL_DIMENSION] = {
            third[0] - first[0],
            third[1] - first[1],
            third[2] - first[2]
        };
        float face_normal[TA_MODEL_DIMENSION] = {
            edge_a[1] * edge_b[2] - edge_a[2] * edge_b[1],
            edge_a[2] * edge_b[0] - edge_a[0] * edge_b[2],
            edge_a[0] * edge_b[1] - edge_a[1] * edge_b[0],
        };
        for (int corner = 0; corner < TA_MODEL_DIMENSION; corner++) {
            size_t position_index = faces[index][corner];
            normal_accum[position_index * TA_MODEL_DIMENSION + 0] += face_normal[0];
            normal_accum[position_index * TA_MODEL_DIMENSION + 1] += face_normal[1];
            normal_accum[position_index * TA_MODEL_DIMENSION + 2] += face_normal[2];
        }
    }

    for (size_t index = 0; index < position_count; index++) {
        float *normal = &normal_accum[index * TA_MODEL_DIMENSION];
        float length  = sqrtf(normal[0] * normal[0] + normal[1] * normal[1] + normal[2] * normal[2]);
        if (length > 0.0f) {
            normal[0] /= length;
            normal[1] /= length;
            normal[2] /= length;
        } else {
            normal[0] = 0.0f;
            normal[1] = 0.0f;
            normal[2] = 1.0f;
        }
    }

    for (size_t index = 0; index < face_count; index++) {
        for (int corner = 0; corner < TA_MODEL_DIMENSION; corner++) {
            size_t position_index = faces[index][corner];
            ta_model_vertex vertex;
            memcpy(vertex.position, positions[position_index], sizeof(vertex.position));
            memcpy(vertex.normal, &normal_accum[position_index * TA_MODEL_DIMENSION], sizeof(vertex.normal));
            if (!ta_model_push_vertex(mesh, &vertex_capacity, &vertex)) {
                goto done;
            }
        }
    }

    success = mesh->vertex_count > 0;
done:
    if (!success) {
        ta_model_free(mesh);
    }
    free(normal_accum);
    return success;
}

void ta_model_free(ta_model_mesh *mesh) {
    free(mesh->vertices);
    mesh->vertices     = NULL;
    mesh->vertex_count = 0;
}

typedef struct ta_model_uniforms {
    float mvp[16];   // TODO: Use dynamic buffer
    float model[16]; // TODO: Use dynamic buffer
    float camera_position[TA_MODEL_DIMENSION];
    float padding;
} ta_model_uniforms;

static SDL_GPUShader *ta_model_shader(SDL_GPUDevice *device, const void *code, size_t size, SDL_GPUShaderStage stage) {
    SDL_GPUShaderCreateInfo info = {
        .code                = code,
        .code_size           = size,
        .entrypoint          = "main",
        .format              = SDL_GPU_SHADERFORMAT_SPIRV,
        .stage               = stage,
        .num_uniform_buffers = stage == SDL_GPU_SHADERSTAGE_VERTEX ? 1 : 0,
    };
    return ta_sdl_create_gpu_shader(device, &info);
}

static void ta_model_identity(float *matrix) {
    for (int index = 0; index < 16; index++) {
        matrix[index] = index % 5 == 0 ? 1.0f : 0.0f;
    }
}

static void ta_model_multiply(float *result, const float *left, const float *right) {
    float product[16];
    for (int column = 0; column < 4; column++) for (int row = 0; row < 4; row++) {
        product[column * 4 + row] = 0.0f;
        for (int inner = 0; inner < 4; inner++) { 
            product[column * 4 + row] += left[inner * 4 + row] * right[column * 4 + inner];
        }
    }
    memcpy(result, product, sizeof(product));
}

bool ta_model_init(ta_model *model, SDL_GPUDevice *device, const float positions[][TA_MODEL_DIMENSION], size_t position_count, const unsigned int faces[][TA_MODEL_DIMENSION], size_t face_count, SDL_GPUTextureFormat color_format, SDL_GPUTextureFormat depth_format, SDL_GPUSampleCount sample_count) {
    memset(model, 0, sizeof(*model));
    if (!ta_model_build(positions, position_count, faces, face_count, &model->mesh) || model->mesh.vertex_count == 0) {
        return false;
    }
    float minimum[TA_MODEL_DIMENSION];
    float maximum[TA_MODEL_DIMENSION];
    memcpy(minimum, model->mesh.vertices[0].position, sizeof(minimum));
    memcpy(maximum, minimum, sizeof(maximum));
    for (size_t index = 1; index < model->mesh.vertex_count; index++) for (int axis = 0; axis < TA_MODEL_DIMENSION; axis++) {
        if (model->mesh.vertices[index].position[axis] < minimum[axis]) {
            minimum[axis] = model->mesh.vertices[index].position[axis];
        }
        if (model->mesh.vertices[index].position[axis] > maximum[axis]) {
            maximum[axis] = model->mesh.vertices[index].position[axis];
        }
    }
    float center[TA_MODEL_DIMENSION] = {
        (minimum[0] + maximum[0]) * 0.5f,
        (minimum[1] + maximum[1]) * 0.5f,
        (minimum[2] + maximum[2]) * 0.5f
    };
    float extent = maximum[0] - minimum[0];
    for (int axis = 1; axis < TA_MODEL_DIMENSION; axis++) {
        if (maximum[axis] - minimum[axis] > extent) {
            extent = maximum[axis] - minimum[axis];
        }
    }
    for (size_t index = 0; index < model->mesh.vertex_count; index++) {
        for (int axis = 0; axis < TA_MODEL_DIMENSION; axis++) {
            model->mesh.vertices[index].position[axis] = (model->mesh.vertices[index].position[axis] - center[axis]) * (2.2f / extent);
        }
    }

    SDL_GPUBufferCreateInfo buffer_info = {
        .usage = SDL_GPU_BUFFERUSAGE_VERTEX,
        .size  = (Uint32)(model->mesh.vertex_count * sizeof(*model->mesh.vertices))
    };
    model->vertex_buffer = SDL_CreateGPUBuffer(device, &buffer_info);
    if (!model->vertex_buffer) {
        ta_model_free(&model->mesh);
        return false;
    }
    model->vertex_count = (Uint32)model->mesh.vertex_count;
    SDL_GPUShader *vertex_shader = ta_model_shader(device, file_model_vert_slang_spv_start, file_model_vert_slang_spv_size, SDL_GPU_SHADERSTAGE_VERTEX);
    SDL_GPUShader *fragment_shader = ta_model_shader(device, file_model_frag_slang_spv_start, file_model_frag_slang_spv_size, SDL_GPU_SHADERSTAGE_FRAGMENT);
    SDL_GPUVertexBufferDescription vertex_buffer = {.slot = 0, .pitch = sizeof(ta_model_vertex), .input_rate = SDL_GPU_VERTEXINPUTRATE_VERTEX};
    SDL_GPUVertexAttribute attributes[2] = {
        {
            .location    = 0,
            .buffer_slot = 0,
            .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3,
            .offset      = 0
        },
        {
            .location    = 1,
            .buffer_slot = 0,
            .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3,
            .offset      = sizeof(float) * TA_MODEL_DIMENSION
        },
    };
    SDL_GPUColorTargetDescription color_target = {
        .format = color_format
    };
    SDL_GPUGraphicsPipelineCreateInfo pipeline_info = {
        .vertex_shader       = vertex_shader,
        .fragment_shader     = fragment_shader,
        .vertex_input_state  = {
            .vertex_buffer_descriptions = &vertex_buffer,
            .num_vertex_buffers         = 1,
            .vertex_attributes          = attributes,
            .num_vertex_attributes      = 2
        },
        .primitive_type      = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
        .rasterizer_state    = {
            .cull_mode = SDL_GPU_CULLMODE_NONE
        },
        .depth_stencil_state = {
            .compare_op         = SDL_GPU_COMPAREOP_LESS,
            .enable_depth_test  = true,
            .enable_depth_write = true
        },
        .multisample_state   = {
            .sample_count = sample_count
        },
        .target_info         = {
            .num_color_targets         = 1,
            .color_target_descriptions = &color_target,
            .depth_stencil_format      = depth_format,
            .has_depth_stencil_target  = true
        },
    };
    model->pipeline = ta_sdl_create_gpu_graphics_pipeline(device, &pipeline_info);
    SDL_ReleaseGPUShader(device, vertex_shader);
    SDL_ReleaseGPUShader(device, fragment_shader);
    return model->pipeline != NULL;
}

void ta_model_upload(ta_model *model, SDL_GPUCommandBuffer *command_buffer, SDL_GPUDevice *device) {
    if (model->uploaded || !model->vertex_buffer) {
        return;
    }
    Uint32 size = (Uint32)(model->mesh.vertex_count * sizeof(*model->mesh.vertices));
    SDL_GPUTransferBufferCreateInfo transfer_info = {
        .usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
        .size  = size
    };
    SDL_GPUTransferBuffer *transfer_buffer = ta_sdl_create_gpu_transfer_buffer(device, &transfer_info);
    void *mapped = ta_sdl_map_gpu_transfer_buffer(device, transfer_buffer, false);
    memcpy(mapped, model->mesh.vertices, size);
    SDL_UnmapGPUTransferBuffer(device, transfer_buffer);
    SDL_GPUCopyPass *copy_pass = ta_sdl_begin_gpu_copy_pass(command_buffer);
    SDL_GPUTransferBufferLocation source = {
        .transfer_buffer = transfer_buffer
    };
    SDL_GPUBufferRegion destination = {
        .buffer = model->vertex_buffer,
        .size   = size
    };
    SDL_UploadToGPUBuffer(copy_pass, &source, &destination, false);
    SDL_EndGPUCopyPass(copy_pass);
    SDL_ReleaseGPUTransferBuffer(device, transfer_buffer);
    ta_model_free(&model->mesh);
    model->uploaded = true;
}

void ta_model_render(const ta_model *model, SDL_GPUCommandBuffer *command_buffer, SDL_GPURenderPass *render_pass, int window_width, int window_height) {
    if (!model->pipeline || !model->vertex_buffer || !model->uploaded) {
        return;
    }
    ta_model_uniforms uniforms;
    float view[16];
    float projection[16] = {0};
    float view_model[16];
    float rotation_x[16];
    float rotation_y[16];
    float rotation_z[16];
    float rotation_xy[16];
    float cosine = cosf(model->rotation[0]);
    float sine = sinf(model->rotation[0]);
    ta_model_identity(rotation_x);
    rotation_x[5] = cosine;
    rotation_x[6] = sine;
    rotation_x[9] = -sine;
    rotation_x[10] = cosine;
    cosine = cosf(model->rotation[1]);
    sine = sinf(model->rotation[1]);
    ta_model_identity(rotation_y);
    rotation_y[0] = cosine;
    rotation_y[2] = -sine;
    rotation_y[8] = sine;
    rotation_y[10] = cosine;
    cosine = cosf(model->rotation[2]);
    sine = sinf(model->rotation[2]);
    ta_model_identity(rotation_z);
    rotation_z[0] = cosine;
    rotation_z[1] = sine;
    rotation_z[4] = -sine;
    rotation_z[5] = cosine;
    ta_model_multiply(rotation_xy, rotation_y, rotation_x);
    ta_model_multiply(uniforms.model, rotation_z, rotation_xy);
    uniforms.model[12] = model->position[0];
    uniforms.model[13] = model->position[1];
    uniforms.model[14] = model->position[2];
    float aspect = (float)window_width / (float)window_height;
    float focal_length = 1.0f / tanf(ta_rad(ta_camera_current.fov * 0.5f));
    float near_plane = 0.1f;
    float far_plane = 100.0f;
    ta_camera_get_view_matrix(view);
    projection[0] = focal_length / aspect;
    projection[5] = focal_length;
    projection[10] = far_plane / (near_plane - far_plane);
    projection[11] = -1.0f;
    projection[14] = near_plane * far_plane / (near_plane - far_plane);
    ta_model_multiply(view_model, view, uniforms.model);
    ta_model_multiply(uniforms.mvp, projection, view_model);
    uniforms.camera_position[0] = ta_camera_current.x;
    uniforms.camera_position[1] = ta_camera_current.y;
    uniforms.camera_position[2] = ta_camera_current.z;
    uniforms.padding = 0.0f;
    SDL_PushGPUVertexUniformData(command_buffer, 0, &uniforms, sizeof(uniforms));
    SDL_BindGPUGraphicsPipeline(render_pass, model->pipeline);
    SDL_GPUBufferBinding binding = {
        .buffer = model->vertex_buffer
    };
    SDL_BindGPUVertexBuffers(render_pass, 0, &binding, 1);
    SDL_DrawGPUPrimitives(render_pass, model->vertex_count, 1, 0, 0);
}

void ta_model_destroy(ta_model *model, SDL_GPUDevice *device) {
    SDL_ReleaseGPUGraphicsPipeline(device, model->pipeline);
    SDL_ReleaseGPUBuffer(device, model->vertex_buffer);
    ta_model_free(&model->mesh);
    memset(model, 0, sizeof(*model));
}