#include "common.h"
#include "world.h"
#include "display.h"
#include "time.h"

#include <inline_c.h>

static CVECTOR colors[] = 
{
    { 255, 0, 0 },
    { 0, 255, 0 },
    { 0, 0, 255 },
    { 255, 255, 0 },
    { 255, 0, 255 },
    { 0, 255, 255 },
    { 128, 255, 0 },
    { 255, 128, 0 },
    { 128, 0, 255 },
    { 255, 0, 128 },
    { 0, 128, 255 },
    { 0, 255, 128 },
};
static const int numColors = sizeof(colors) / sizeof(CVECTOR);

// Set data pointers directly from an in-memory byte buffer
#define LOAD_CHUNK(type, dst, num, src, entry)   \
    dst = (type*)((src) + (entry).offset);  \
    num = (entry).size / sizeof(type);

// Allocate memory per chunk and copy data from a byte buffer
// #define LOAD_CHUNK(type, dst, num, src, entry)  \
//     dst = (type*)malloc((entry).size);  \
//     memcpy(dst, (src) + (entry).offset, (entry).size);  \
//     num = (entry).size / sizeof(type);

void world_load(const u_long *data, world_t *world)
{
    const char *bytes = (const char*)data;

    ps1bsp_header_t* header = (ps1bsp_header_t*)bytes;

    LOAD_CHUNK(ps1bsp_texture_t, world->textures, world->numTextures, bytes, header->textures);
    LOAD_CHUNK(ps1bsp_vertex_t, world->vertices, world->numVertices, bytes, header->vertices);
    LOAD_CHUNK(ps1bsp_face_t, world->faces, world->numFaces, bytes, header->faces);
    LOAD_CHUNK(ps1bsp_facevertex_t, world->faceVertices, world->numFaceVertices, bytes, header->faceVertices);
    LOAD_CHUNK(ps1bsp_plane_t, world->planes, world->numPlanes, bytes, header->planes);
    LOAD_CHUNK(ps1bsp_node_t, world->nodes, world->numNodes, bytes, header->nodes);
    LOAD_CHUNK(ps1bsp_leaf_t, world->leaves, world->numLeaves, bytes, header->leaves);
    LOAD_CHUNK(u_short, world->leafFaces, world->numLeafFaces, bytes, header->leafFaces);
    LOAD_CHUNK(u_char, world->visData, world->numVisData, bytes, header->visData);
}

static INLINE void drawface_triangle_fan(const ps1bsp_face_t *face, SVECTOR *vecs)
{
    int p;

    if (!mem_checkprim(sizeof(POLY_G3), face->numFaceVertices - 2))
        return;

    // Draw the face as a triangle fan
    u_char maxVert = face->numFaceVertices - 1;
    for (int vertIdx = 1; vertIdx < maxVert; ++vertIdx)
    {
        const SVECTOR *v0 = &vecs[0];
        const SVECTOR *v1 = &vecs[vertIdx];
        const SVECTOR *v2 = &vecs[vertIdx + 1];

        // Naively draw the triangle with GTE, nothing special or optimized about this
        gte_ldv3(v0, v1, v2);
        gte_rtpt();     // Rotation, translation, perspective projection

        // Average Z for depth sorting and culling
        gte_avsz3();
        gte_stotz(&p);
        short depth = p >> 2;
        if (depth <= 0 || depth >= OTLEN)
            continue;

        // Draw a flat-shaded untextured colored triangle
        POLY_G3 *poly = (POLY_G3*)mem_prim(sizeof(POLY_G3));
        setPolyG3(poly);
        gte_stsxy3_g3(poly);

        poly->r0 = poly->g0 = poly->b0 = (uint8_t)v0->pad;
        poly->r1 = poly->g1 = poly->b1 = (uint8_t)v1->pad;
        poly->r2 = poly->g2 = poly->b2 = (uint8_t)v2->pad;

        addPrim(curOT + depth, poly);
        ++polyCount;
    }
}

static INLINE void drawface_triangle_strip(const ps1bsp_face_t *face, SVECTOR *vecs)
{
    int p;

    u_char numTris = face->numFaceVertices - 2;
    if (!mem_checkprim(sizeof(POLY_G3), numTris))
        return;

    // Draw the face as a triangle strip
    const SVECTOR *v0, *v1, *v2;
    const SVECTOR *head = vecs;
    const SVECTOR *tail = vecs + face->numFaceVertices;
    u_char reverse = 0;

    v2 = head++;    // Initialize first vertex to index 0 and set head to index 1

    for (u_char triIdx = 0; triIdx < numTris; ++triIdx)
    {
        if (reverse ^= 1)
        {
            v0 = v2;
            v1 = head;
            v2 = --tail;
        }
        else
        {
            v0 = v1;
            v1 = ++head;
            v2 = tail;
        }

        // Naively draw the triangle with GTE, nothing special or optimized about this
        gte_ldv3(v0, v1, v2);
        gte_rtpt();     // Rotation, translation, perspective projection

        // Average Z for depth sorting and culling
        gte_avsz3();
        gte_stotz(&p);
        short depth = p >> 2;
        if (depth <= 0 || depth >= OTLEN)
            continue;

        // Draw a flat-shaded untextured colored triangle
        POLY_G3 *poly = (POLY_G3*)mem_prim(sizeof(POLY_G3));
        setPolyG3(poly);
        gte_stsxy3_g3(poly);

        poly->r0 = poly->g0 = poly->b0 = (uint8_t)v0->pad;
        poly->r1 = poly->g1 = poly->b1 = (uint8_t)v1->pad;
        poly->r2 = poly->g2 = poly->b2 = (uint8_t)v2->pad;

        addPrim(curOT + depth, poly);
        ++polyCount;
    }
}

static INLINE void drawface_quad_strip(const ps1bsp_face_t *face, SVECTOR *vecs)
{
    int p;

    u_char numQuads = (face->numFaceVertices - 1) / 2;
    if (!mem_checkprim(sizeof(POLY_G4), numQuads))
        return;

    // Draw the face as a quad strip
    const SVECTOR *v0, *v1, *v2, *v3;
    const SVECTOR *head = vecs;
    const SVECTOR *tail = vecs + face->numFaceVertices;

    // Initialize the first two vertices
    v2 = --tail;
    v3 = head++;

    // Normally a quad strip would have (N-2)/2 quads, but we might end up with a sole triangle at the end which will be drawn as a collapsed quad
    for (u_char quadIdx = 0; quadIdx < numQuads; ++quadIdx)
    {
        v0 = v2;
        v1 = v3;
        v2 = --tail;
        v3 = head++;

        // Naively draw the quad with GTE, nothing special or optimized about this
        gte_ldv3(v0, v1, v2);
        gte_rtpt();     // Rotation, translation, perspective projection

        // Average Z for depth sorting and culling
        gte_avsz3();
        gte_stotz(&p);
        short depth = p >> 2;
        if (depth <= 0 || depth >= OTLEN)
            continue;

        // Draw a flat-shaded untextured colored quad
        POLY_G4 *poly = (POLY_G4*)mem_prim(sizeof(POLY_G4));
        setPolyG4(poly);
        gte_stsxy0(&poly->x0);
        gte_stsxy1(&poly->x1);
        gte_stsxy2(&poly->x2);

        // Transform the fourth vertex to complete the quad
        gte_ldv0(v3);
        gte_rtps();
        gte_stsxy(&poly->x3);

        poly->r0 = poly->g0 = poly->b0 = (uint8_t)v0->pad;
        poly->r1 = poly->g1 = poly->b1 = (uint8_t)v1->pad;
        poly->r2 = poly->g2 = poly->b2 = (uint8_t)v2->pad;
        poly->r3 = poly->g3 = poly->b3 = (uint8_t)v3->pad;

        addPrim(curOT + depth, poly);
        ++polyCount;
    }
}

static INLINE void drawface_quad_strip_tex(const ps1bsp_face_t *face, STVECTOR *vecs, u_short tpage)
{
    int p;

    u_char numQuads = (face->numFaceVertices - 1) / 2;
    if (!mem_checkprim(sizeof(POLY_GT4), numQuads))
        return;

    // Draw the face as a quad strip
    const STVECTOR *v0, *v1, *v2, *v3;
    const STVECTOR *head = vecs;
    const STVECTOR *tail = vecs + face->numFaceVertices;

    // Initialize the first two vertices
    v2 = --tail;
    v3 = head++;

    // Normally a quad strip would have (N-2)/2 quads, but we might end up with a sole triangle at the end which will be drawn as a collapsed quad
    for (u_char quadIdx = 0; quadIdx < numQuads; ++quadIdx)
    {
        v0 = v2;
        v1 = v3;
        v2 = --tail;
        v3 = head++;

        // Naively draw the quad with GTE, nothing special or optimized about this
        gte_ldv3(v0, v1, v2);
        gte_rtpt();     // Rotation, translation, perspective projection

        // Average Z for depth sorting and culling
        gte_avsz3();
        gte_stotz(&p);
        short depth = p >> 2;
        if (depth <= 0 || depth >= OTLEN)
            continue;

        // Draw a flat-shaded untextured colored quad
        POLY_GT4 *poly = (POLY_GT4*)mem_prim(sizeof(POLY_GT4));
        setPolyGT4(poly);
        gte_stsxy0(&poly->x0);
        gte_stsxy1(&poly->x1);
        gte_stsxy2(&poly->x2);

        // Transform the fourth vertex to complete the quad
        gte_ldv0(v3);
        gte_rtps();
        gte_stsxy(&poly->x3);

        // Texture UVs
        setUV4(poly, v0->u, v0->v, v1->u, v1->v, v2->u, v2->v, v3->u, v3->v);
        poly->clut = quake_clut;
        poly->tpage = tpage;

        // Vertex color lighting
        poly->r0 = poly->g0 = poly->b0 = (uint8_t)v0->pad;
        poly->r1 = poly->g1 = poly->b1 = (uint8_t)v1->pad;
        poly->r2 = poly->g2 = poly->b2 = (uint8_t)v2->pad;
        poly->r3 = poly->g3 = poly->b3 = (uint8_t)v3->pad;

        addPrim(curOT + depth, poly);
        ++polyCount;
    }
}

static INLINE short world_pointPlaneDist(const VECTOR *point, const ps1bsp_plane_t *plane)
{
    // Make use of axis-aligned planes to skip the need for a dot product
    if (plane->type < 3)
        return (short)(((int*)point)[plane->type] - plane->dist);
    
    return (short)m_pointPlaneDist2(*point, plane->normal, plane->dist);
}

static INLINE short world_planeDot(const SVECTOR *dir, const ps1bsp_plane_t *plane)
{
    // Make use of axis-aligned planes to skip the need for a dot product
    if (plane->type < 3)
        return ((short*)dir)[plane->type];

    return (short)m_dot12s(*dir, plane->normal);
}

static void world_drawface(const world_t *world, const ps1bsp_face_t *face)
{
    // TODO: do an early primitive buffer check here, so we can skip backface culling & vertex copying if it's already full

    // Backface culling using the face's plane and center point
    // This eliminates the need for normal clipping per polygon
    SVECTOR cam_vec;
    cam_vec.vx = face->center.vx - cam_pos.vx;
    cam_vec.vy = face->center.vy - cam_pos.vy;
    cam_vec.vz = face->center.vz - cam_pos.vz;
    
    const ps1bsp_plane_t *plane = &world->planes[face->planeId];
    // NOTE: this value could be REALLY useful for determining the tessellation subdivisions. It has camera distance *and* angle in it.
    // Just include the face size/area for an approximate screen size. Maybe also separate x/y/z for angle-dependent tessellation.
    short dot = world_planeDot(&cam_vec, plane);
    if ((dot >= 0) ^ face->side)
        return;

    if (!enableTexturing)
    {
        // Draw untextured, vertex colored polygons
        SVECTOR *vecs = (SVECTOR*)(scratchpad + 256);

        // Copy this face's vertices into scratch RAM for fast reuse
        // TODO: this is the main performance bottleneck right now!
        ps1bsp_facevertex_t *faceVertex = &world->faceVertices[face->firstFaceVertex];
        SVECTOR *curVec = vecs;
        for (int vertIdx = 0; vertIdx < face->numFaceVertices; ++vertIdx, ++faceVertex, ++curVec)
        {
            const ps1bsp_vertex_t *vert = &world->vertices[faceVertex->index];
            *curVec = *((SVECTOR*)vert);
            curVec->pad = faceVertex->light;
        }

        if (face->numFaceVertices == 3) // Special case: draw single triangles using the simplest method
            drawface_triangle_fan(face, vecs);
        else
            drawface_quad_strip(face, vecs);
    }
    else
    {
        // Draw textured, vertex colored polygons
        STVECTOR *vecs = (STVECTOR*)(scratchpad + 256);

        // Copy this face's vertices into scratch RAM for fast reuse
        ps1bsp_facevertex_t *faceVertex = &world->faceVertices[face->firstFaceVertex];
        ps1bsp_texture_t *faceTexture = &world->textures[face->textureId];
        STVECTOR *curVec = vecs;
        for (int vertIdx = 0; vertIdx < face->numFaceVertices; ++vertIdx, ++faceVertex, ++curVec)
        {
            const ps1bsp_vertex_t *vert = &world->vertices[faceVertex->index];
            *((SVECTOR*)curVec) = *((SVECTOR*)vert);
            curVec->u = (u_short)faceVertex->u;
            curVec->v = (u_short)faceVertex->v;
            curVec->pad = faceVertex->light;
        }

        drawface_quad_strip_tex(face, vecs, faceTexture->tpage);
    }
}

static void world_drawnode(const world_t *world, short nodeIdx, u_char *pvs)
{
    u_long frameNum = time_getFrameNumber();

    if (nodeIdx < 0)    // Leaf node
    {
        // Check if this leaf is visible from the current camera position
        u_short test = ~nodeIdx - 1;
        if ((pvs[test >> 3] & (1 << (test & 0x7))) == 0)
            return;

        const ps1bsp_leaf_t *leaf = &world->leaves[~nodeIdx];
        const u_short *leafFace = &world->leafFaces[leaf->firstLeafFace];

        for (u_short leafFaceIdx = 0; leafFaceIdx < leaf->numLeafFaces; ++leafFaceIdx, ++leafFace)
        {
            ps1bsp_face_t *face = &world->faces[*leafFace];

            // Check if we've already drawn this face on the current frame
            if (face->drawFrame == frameNum)
                continue;

            world_drawface(world, face);
            face->drawFrame = frameNum;
        }

        return;
    }

    const ps1bsp_node_t *node = &world->nodes[nodeIdx];
    const ps1bsp_plane_t *plane = &world->planes[node->planeId];
    short dist = world_pointPlaneDist(&cam_pos, plane);
    
    // Draw child nodes in front-to-back order; adding faces to the OT will reverse the drawing order
    char order = dist < 0;
    world_drawnode(world, node->children[order], pvs);
    world_drawnode(world, node->children[order ^ 1], pvs);
}

// Decompress PVS data for the given leaf ID and store it in RAM at the given buffer pointer location.
// Returns the memory location of decompressed PVS data, and moves the buffer pointer forward.
static u_char *world_loadVisData(const world_t *world, u_short leafIdx, u_char **buffer)
{
    u_char *head = *buffer;
    u_char *tail = head;

    const ps1bsp_leaf_t *leaf = &world->leaves[leafIdx];
    
    const u_char *v = &world->visData[leaf->vislist];
    for (int l = 1; l < world->numLeaves; )
    {
        u_char bits = *v++;
        if (bits)
        {
            *tail++ = bits;
            l += 8;
        }
        else
        {
            u_char skip = *v++;
            for (u_char i = 0; i < skip; ++i, l += 8)
            {
                *tail++ = 0;
            }
        }
    }

    *buffer = tail;
    return head;
}

static u_char *world_noVisData(const world_t *world, u_char **buffer)
{
    u_char *head = *buffer;
    u_char *tail = head;

    for (int l = 1; l < world->numLeaves; l += 8)
    {
        *tail++ = 0xFF;
    }

    *buffer = tail;
    return head;
}

static u_short world_leafAtPoint(const world_t *world, const VECTOR *point)
{
    short nodeIdx = 0;
    while (nodeIdx >= 0)
    {
        const ps1bsp_node_t *node = &world->nodes[nodeIdx];
        const ps1bsp_plane_t *plane = &world->planes[node->planeId];

        short dist = world_pointPlaneDist(point, plane);
        nodeIdx = node->children[(dist > 0) ^ 1];
    }

    return ~nodeIdx;
}

void world_draw(const world_t *world)
{
    int p;

    // The world doesn't move, so we just set the camera view-projection matrix
    gte_SetRotMatrix(&vp_matrix);
	gte_SetTransMatrix(&vp_matrix);

    cam_leaf = world_leafAtPoint(world, &cam_pos);

    u_char *pvsbuf = scratchpad;
    u_char *pvs = world_loadVisData(world, cam_leaf, &pvsbuf);
    //u_char *pvs = world_noVisData(world, &pvsbuf);

    world_drawnode(world, 0, pvs);
}