#include "demo3d_bsp_render.h"

#include <string.h>

#include "pxl8_gfx3d.h"
#include "pxl8_log.h"
#include "pxl8_mem.h"
#include "pxl8_mesh.h"

static inline bool demo3d_bsp_get_edge_vertex(const demo3d_bsp* bsp, i32 surfedge_idx, u32* out_vert_idx) {
    if (surfedge_idx >= (i32)bsp->num_surfedges) return false;

    i32 edge_idx = bsp->surfedges[surfedge_idx];
    u32 vertex_index;

    if (edge_idx >= 0) {
        if ((u32)edge_idx >= bsp->num_edges) return false;
        vertex_index = 0;
    } else {
        edge_idx = -edge_idx;
        if ((u32)edge_idx >= bsp->num_edges) return false;
        vertex_index = 1;
    }

    *out_vert_idx = bsp->edges[edge_idx].vertex[vertex_index];
    return *out_vert_idx < bsp->num_vertices;
}

static inline bool face_in_frustum(const demo3d_bsp* bsp, u32 face_id, const pxl8_frustum* frustum) {
    const demo3d_bsp_face* face = &bsp->faces[face_id];
    return pxl8_frustum_test_aabb(frustum, face->aabb_min, face->aabb_max);
}

static void collect_face_to_mesh(const demo3d_bsp* bsp, const demo3d_bsp_render_state* state,
                                 u32 face_id, pxl8_mesh* mesh, u8 ambient) {
    const demo3d_bsp_face* face = &bsp->faces[face_id];
    if (face->num_edges < 3) return;

    pxl8_vec3 normal = {0, 1, 0};
    if (face->plane_id < bsp->num_planes) {
        normal = bsp->planes[face->plane_id].normal;
        if (face->side) {
            normal.x = -normal.x;
            normal.y = -normal.y;
            normal.z = -normal.z;
        }
    }

    const pxl8_gfx_material* material = NULL;
    f32 tex_scale = 64.0f;
    if (state && face->material_id < state->num_materials) {
        material = &state->materials[face->material_id];
    }

    pxl8_vec3 u_axis = {1.0f, 0.0f, 0.0f};
    pxl8_vec3 v_axis = {0.0f, 0.0f, 1.0f};
    f32 u_offset = 0.0f, v_offset = 0.0f;
    if (material) {
        u_offset = material->u_offset;
        v_offset = material->v_offset;
    }
    f32 abs_ny = normal.y < 0 ? -normal.y : normal.y;
    if (abs_ny > 0.7f) {
        u_axis = (pxl8_vec3){1.0f, 0.0f, 0.0f};
        v_axis = (pxl8_vec3){0.0f, 0.0f, 1.0f};
    } else {
        f32 abs_nx = normal.x < 0 ? -normal.x : normal.x;
        f32 abs_nz = normal.z < 0 ? -normal.z : normal.z;
        if (abs_nx > abs_nz) {
            u_axis = (pxl8_vec3){0.0f, 0.0f, 1.0f};
        } else {
            u_axis = (pxl8_vec3){1.0f, 0.0f, 0.0f};
        }
        v_axis = (pxl8_vec3){0.0f, 1.0f, 0.0f};
    }

    u16 base_idx = (u16)mesh->vertex_count;
    u32 num_verts = 0;

    for (u32 i = 0; i < face->num_edges && num_verts < 64; i++) {
        u32 edge_i = face->side ? (face->num_edges - 1 - i) : i;
        i32 surfedge_idx = face->first_edge + edge_i;
        u32 vert_idx;

        if (!demo3d_bsp_get_edge_vertex(bsp, surfedge_idx, &vert_idx)) {
            continue;
        }

        pxl8_vec3 pos = bsp->vertices[vert_idx].position;

        f32 u = (pxl8_vec3_dot(pos, u_axis) + u_offset) / tex_scale;
        f32 v = (pxl8_vec3_dot(pos, v_axis) + v_offset) / tex_scale;

        u8 light = 255;
        if (bsp->vertex_lights && vert_idx < bsp->num_vertex_lights) {
            u32 packed = bsp->vertex_lights[vert_idx];
            u8 direct = (packed >> 24) & 0xFF;
            u8 ao = (packed >> 16) & 0xFF;
            f32 combined = (f32)direct + ((f32)ambient / 255.0f) * (f32)ao;
            light = (u8)(combined > 255.0f ? 255.0f : combined);
        }

        pxl8_vertex vtx = {
            .position = pos,
            .normal = normal,
            .u = u,
            .v = v,
            .color = 15,
            .light = light,
        };
        pxl8_mesh_push_vertex(mesh, vtx);
        num_verts++;
    }

    if (num_verts < 3) return;

    for (u32 i = 1; i < num_verts - 1; i++) {
        pxl8_mesh_push_triangle(mesh, base_idx, base_idx + i, base_idx + i + 1);
    }
}

demo3d_bsp_render_state* demo3d_bsp_render_state_create(u32 num_faces) {
    demo3d_bsp_render_state* state = pxl8_calloc(1, sizeof(demo3d_bsp_render_state));
    if (!state) return NULL;

    state->num_faces = num_faces;
    if (num_faces > 0) {
        state->render_face_flags = pxl8_calloc(num_faces, 1);
        if (!state->render_face_flags) {
            pxl8_free(state);
            return NULL;
        }
    }

    return state;
}

void demo3d_bsp_render_state_destroy(demo3d_bsp_render_state* state) {
    if (!state) return;
    pxl8_free(state->materials);
    pxl8_free(state->render_face_flags);
    pxl8_free(state);
}

void demo3d_bsp_render(pxl8_gfx* gfx, const demo3d_bsp* bsp,
                     demo3d_bsp_render_state* state,
                     const pxl8_gfx_draw_opts* opts) {
    if (!gfx || !bsp || !state || bsp->num_faces == 0) return;
    if (!state->materials || state->num_materials == 0) return;
    if (!state->render_face_flags) return;

    const pxl8_frustum* frustum = pxl8_3d_get_frustum(gfx);
    if (!frustum) return;

    pxl8_mesh* mesh = pxl8_mesh_create(8192, 16384);
    if (!mesh) return;

    u8 ambient = pxl8_gfx_get_ambient(gfx);
    pxl8_mat4 identity = pxl8_mat4_identity();

    for (u32 mat = 0; mat < state->num_materials; mat++) {
        for (u32 face_id = 0; face_id < bsp->num_faces; face_id++) {
            if (!state->render_face_flags[face_id]) continue;
            const demo3d_bsp_face* face = &bsp->faces[face_id];
            if (face->material_id != mat) continue;
            if (!face_in_frustum(bsp, face_id, frustum)) continue;
            collect_face_to_mesh(bsp, state, face_id, mesh, ambient);
        }
        if (mesh->index_count > 0) {
            pxl8_gfx_material mat_copy = state->materials[mat];
            if (state->exterior) mat_copy.double_sided = true;
            pxl8_3d_draw_mesh(gfx, mesh, &identity, &mat_copy, opts);
            pxl8_mesh_clear(mesh);
        }
    }

    pxl8_mesh_destroy(mesh);
}

void demo3d_bsp_set_material(demo3d_bsp_render_state* state, u16 material_id, const pxl8_gfx_material* material) {
    if (!state || !material) return;

    if (material_id >= state->num_materials) {
        u32 new_count = material_id + 1;
        pxl8_gfx_material* new_materials = pxl8_realloc(state->materials, new_count * sizeof(pxl8_gfx_material));
        if (!new_materials) return;

        for (u32 i = state->num_materials; i < new_count; i++) {
            memset(&new_materials[i], 0, sizeof(pxl8_gfx_material));
            new_materials[i].u_axis = (pxl8_vec3){1.0f, 0.0f, 0.0f};
            if (i == 0 || i == 2) {
                new_materials[i].v_axis = (pxl8_vec3){0.0f, 0.0f, 1.0f};
            } else {
                new_materials[i].v_axis = (pxl8_vec3){0.0f, 1.0f, 0.0f};
            }
        }

        state->materials = new_materials;
        state->num_materials = new_count;
    }

    pxl8_vec3 u_axis = state->materials[material_id].u_axis;
    pxl8_vec3 v_axis = state->materials[material_id].v_axis;
    f32 u_offset = state->materials[material_id].u_offset;
    f32 v_offset = state->materials[material_id].v_offset;

    state->materials[material_id] = *material;

    state->materials[material_id].u_axis = u_axis;
    state->materials[material_id].v_axis = v_axis;
    state->materials[material_id].u_offset = u_offset;
    state->materials[material_id].v_offset = v_offset;
}