#include "pxl8_bsp_render.h"

#include <string.h>

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

typedef struct {
    f32 x0, y0, x1, y1;
} screen_rect;

typedef struct {
    u32 leaf;
    screen_rect window;
} portal_queue_entry;

static inline bool pxl8_bsp_get_edge_vertex(const pxl8_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 pxl8_bsp* bsp, u32 face_id, const pxl8_frustum* frustum) {
    const pxl8_bsp_face* face = &bsp->faces[face_id];
    return pxl8_frustum_test_aabb(frustum, face->aabb_min, face->aabb_max);
}

static inline bool screen_rect_valid(screen_rect r) {
    return r.x0 < r.x1 && r.y0 < r.y1;
}

static inline screen_rect screen_rect_intersect(screen_rect a, screen_rect b) {
    return (screen_rect){
        .x0 = (a.x0 > b.x0) ? a.x0 : b.x0,
        .y0 = (a.y0 > b.y0) ? a.y0 : b.y0,
        .x1 = (a.x1 < b.x1) ? a.x1 : b.x1,
        .y1 = (a.y1 < b.y1) ? a.y1 : b.y1,
    };
}

static inline void expand_rect_with_ndc(screen_rect* r, f32 nx, f32 ny) {
    if (nx < r->x0) r->x0 = nx;
    if (nx > r->x1) r->x1 = nx;
    if (ny < r->y0) r->y0 = ny;
    if (ny > r->y1) r->y1 = ny;
}

static screen_rect project_portal_to_screen(const pxl8_bsp_portal* portal, const pxl8_mat4* vp, f32 wall_height) {
    pxl8_vec3 world_corners[4] = {
        {portal->x0, 0, portal->z0},
        {portal->x1, 0, portal->z1},
        {portal->x1, wall_height, portal->z1},
        {portal->x0, wall_height, portal->z0},
    };

    pxl8_vec4 clip[4];
    bool in_front[4];
    i32 front_count = 0;

    const f32 NEAR_W = 0.001f;

    for (i32 i = 0; i < 4; i++) {
        clip[i] = pxl8_mat4_multiply_vec4(*vp, (pxl8_vec4){world_corners[i].x, world_corners[i].y, world_corners[i].z, 1.0f});
        in_front[i] = clip[i].w > NEAR_W;
        if (in_front[i]) front_count++;
    }

    if (front_count == 0) return (screen_rect){0, 0, 0, 0};

    screen_rect result = {1e30f, 1e30f, -1e30f, -1e30f};

    for (i32 i = 0; i < 4; i++) {
        i32 j = (i + 1) % 4;
        pxl8_vec4 a = clip[i];
        pxl8_vec4 b = clip[j];
        bool a_in = in_front[i];
        bool b_in = in_front[j];

        if (a_in) {
            f32 inv_w = 1.0f / a.w;
            expand_rect_with_ndc(&result, a.x * inv_w, a.y * inv_w);
        }

        if (a_in != b_in) {
            f32 t = (NEAR_W - a.w) / (b.w - a.w);
            pxl8_vec4 intersection = {
                a.x + t * (b.x - a.x),
                a.y + t * (b.y - a.y),
                a.z + t * (b.z - a.z),
                NEAR_W
            };
            f32 inv_w = 1.0f / intersection.w;
            expand_rect_with_ndc(&result, intersection.x * inv_w, intersection.y * inv_w);
        }
    }

    if (result.x0 < -1.0f) result.x0 = -1.0f;
    if (result.y0 < -1.0f) result.y0 = -1.0f;
    if (result.x1 > 1.0f) result.x1 = 1.0f;
    if (result.y1 > 1.0f) result.y1 = 1.0f;

    return result;
}

static void collect_face_to_mesh(const pxl8_bsp* bsp, const pxl8_bsp_render_state* state,
                                 u32 face_id, pxl8_mesh* mesh, u8 ambient) {
    const pxl8_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 (!pxl8_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);
    }
}

pxl8_bsp_render_state* pxl8_bsp_render_state_create(u32 num_faces) {
    pxl8_bsp_render_state* state = pxl8_calloc(1, sizeof(pxl8_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 pxl8_bsp_render_state_destroy(pxl8_bsp_render_state* state) {
    if (!state) return;
    pxl8_free(state->materials);
    pxl8_free(state->render_face_flags);
    pxl8_free(state);
}

void pxl8_bsp_render_face(pxl8_gfx* gfx, const pxl8_bsp* bsp, u32 face_id, const pxl8_gfx_material* material) {
    if (!gfx || !bsp || face_id >= bsp->num_faces || !material) return;

    pxl8_mesh* mesh = pxl8_mesh_create(64, 192);
    if (!mesh) return;

    collect_face_to_mesh(bsp, NULL, face_id, mesh, pxl8_gfx_get_ambient(gfx));

    if (mesh->index_count > 0) {
        pxl8_mat4 identity = pxl8_mat4_identity();
        pxl8_3d_draw_mesh(gfx, mesh, &identity, material);
    }

    pxl8_mesh_destroy(mesh);
}

void pxl8_bsp_render(pxl8_gfx* gfx, const pxl8_bsp* bsp,
                     pxl8_bsp_render_state* state, pxl8_vec3 camera_pos) {
    if (!gfx || !bsp || !state || bsp->num_faces == 0) return;
    if (!state->materials || state->num_materials == 0) return;

    const pxl8_frustum* frustum = pxl8_3d_get_frustum(gfx);
    const pxl8_mat4* vp = pxl8_3d_get_view_proj(gfx);
    if (!frustum || !vp) return;

    i32 camera_leaf = pxl8_bsp_find_leaf(bsp, camera_pos);
    if (camera_leaf < 0 || (u32)camera_leaf >= bsp->num_leafs) return;

    if (!bsp->cell_portals || bsp->num_cell_portals == 0) {
        pxl8_mesh* mesh = pxl8_mesh_create(8192, 16384);
        if (!mesh) return;

        u32 current_material = 0xFFFFFFFF;

        for (u32 face_id = 0; face_id < bsp->num_faces; face_id++) {
            if (!face_in_frustum(bsp, face_id, frustum)) continue;

            const pxl8_bsp_face* face = &bsp->faces[face_id];
            u32 material_id = face->material_id;
            if (material_id >= state->num_materials) continue;

            if (material_id != current_material && mesh->index_count > 0 && current_material < state->num_materials) {
                pxl8_mat4 identity = pxl8_mat4_identity();
                pxl8_3d_draw_mesh(gfx, mesh, &identity, &state->materials[current_material]);
                pxl8_mesh_clear(mesh);
            }

            current_material = material_id;
            collect_face_to_mesh(bsp, state, face_id, mesh, pxl8_gfx_get_ambient(gfx));
        }

        if (mesh->index_count > 0 && current_material < state->num_materials) {
            pxl8_mat4 identity = pxl8_mat4_identity();
            pxl8_3d_draw_mesh(gfx, mesh, &identity, &state->materials[current_material]);
        }

        pxl8_mesh_destroy(mesh);
        return;
    }

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

    pxl8_bsp_pvs pvs = pxl8_bsp_decompress_pvs(bsp, camera_leaf);

    u32 visited_bytes = (bsp->num_leafs + 7) / 8;
    u8* visited = pxl8_calloc(visited_bytes, 1);
    screen_rect* cell_windows = pxl8_calloc(bsp->num_leafs, sizeof(screen_rect));
    portal_queue_entry* queue = pxl8_malloc(bsp->num_leafs * 4 * sizeof(portal_queue_entry));
    if (!visited || !cell_windows || !queue) {
        pxl8_free(visited);
        pxl8_free(cell_windows);
        pxl8_free(queue);
        pxl8_bsp_pvs_destroy(&pvs);
        return;
    }

    u32 head = 0, tail = 0;
    screen_rect full_screen = {-1.0f, -1.0f, 1.0f, 1.0f};

    visited[camera_leaf >> 3] |= (1 << (camera_leaf & 7));
    cell_windows[camera_leaf] = full_screen;
    queue[tail++] = (portal_queue_entry){camera_leaf, full_screen};

    f32 wall_height = 128.0f;

    while (head < tail) {
        portal_queue_entry entry = queue[head++];
        u32 leaf_id = entry.leaf;
        screen_rect window = entry.window;

        if (leaf_id >= bsp->num_cell_portals) continue;
        const pxl8_bsp_cell_portals* cp = &bsp->cell_portals[leaf_id];

        for (u8 i = 0; i < cp->num_portals; i++) {
            const pxl8_bsp_portal* portal = &cp->portals[i];
            u32 target = portal->target_leaf;

            if (target >= bsp->num_leafs) continue;
            if (bsp->leafs[target].contents == -1) continue;
            if (!pxl8_bsp_pvs_is_visible(&pvs, target)) continue;

            screen_rect portal_rect = project_portal_to_screen(portal, vp, wall_height);
            if (!screen_rect_valid(portal_rect)) continue;

            screen_rect new_window = screen_rect_intersect(window, portal_rect);
            if (!screen_rect_valid(new_window)) continue;

            u32 byte = target >> 3;
            u32 bit = target & 7;
            if (visited[byte] & (1 << bit)) {
                screen_rect existing = cell_windows[target];
                bool expanded = false;
                if (new_window.x0 < existing.x0) { cell_windows[target].x0 = new_window.x0; expanded = true; }
                if (new_window.y0 < existing.y0) { cell_windows[target].y0 = new_window.y0; expanded = true; }
                if (new_window.x1 > existing.x1) { cell_windows[target].x1 = new_window.x1; expanded = true; }
                if (new_window.y1 > existing.y1) { cell_windows[target].y1 = new_window.y1; expanded = true; }
                if (expanded && tail < bsp->num_leafs * 4) {
                    queue[tail++] = (portal_queue_entry){target, cell_windows[target]};
                }
                continue;
            }

            visited[byte] |= (1 << bit);
            cell_windows[target] = new_window;
            queue[tail++] = (portal_queue_entry){target, new_window};
        }
    }

    pxl8_mesh* mesh = pxl8_mesh_create(8192, 16384);
    if (!mesh) {
        pxl8_free(visited);
        pxl8_free(cell_windows);
        pxl8_free(queue);
        return;
    }

    u32 current_material = 0xFFFFFFFF;

    for (u32 leaf_id = 0; leaf_id < bsp->num_leafs; leaf_id++) {
        if (!(visited[leaf_id >> 3] & (1 << (leaf_id & 7)))) continue;
        if (bsp->leafs[leaf_id].contents == -1) continue;

        const pxl8_bsp_leaf* leaf = &bsp->leafs[leaf_id];

        for (u32 i = 0; i < leaf->num_marksurfaces; i++) {
            u32 surf_idx = leaf->first_marksurface + i;
            if (surf_idx >= bsp->num_marksurfaces) continue;

            u32 face_id = bsp->marksurfaces[surf_idx];
            if (face_id >= bsp->num_faces) continue;

            if (state->render_face_flags[face_id]) continue;
            state->render_face_flags[face_id] = 1;

            if (!face_in_frustum(bsp, face_id, frustum)) continue;

            const pxl8_bsp_face* face = &bsp->faces[face_id];
            u32 material_id = face->material_id;
            if (material_id >= state->num_materials) continue;

            if (material_id != current_material && mesh->index_count > 0 && current_material < state->num_materials) {
                pxl8_mat4 identity = pxl8_mat4_identity();
                pxl8_3d_draw_mesh(gfx, mesh, &identity, &state->materials[current_material]);
                pxl8_mesh_clear(mesh);
            }

            current_material = material_id;
            collect_face_to_mesh(bsp, state, face_id, mesh, pxl8_gfx_get_ambient(gfx));
        }
    }

    if (mesh->index_count > 0 && current_material < state->num_materials) {
        pxl8_mat4 identity = pxl8_mat4_identity();
        pxl8_3d_draw_mesh(gfx, mesh, &identity, &state->materials[current_material]);
    }

    pxl8_bsp_pvs_destroy(&pvs);
    pxl8_free(visited);
    pxl8_free(cell_windows);
    pxl8_free(queue);
    pxl8_mesh_destroy(mesh);
}

void pxl8_bsp_set_material(pxl8_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;
}