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

static P_COLOR colors[] = 
{
    { (255 << 0), 0 },  // Red
    { (255 << 8), 0 },  // Green
    { (255 << 16), 0 }, // Blue
    { (255 << 0) | (255 << 8), 0 },
    { (255 << 0) | (255 << 16), 0 },
    { (255 << 8) | (255 << 16), 0 },
    { (128 << 8) | (255 << 8), 0 },
    { (255 << 0) | (128 << 8), 0 },
    { (128 << 0) | (255 << 16), 0 },
    { (255 << 0) | (128 << 16), 0 },
    { (128 << 8) | (255 << 16), 0 },
    { (128 << 8) | (128 << 16), 0 },
};
static const int numColors = sizeof(colors) / sizeof(P_COLOR);

// 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_worldspawn_t, world->worldSpawn, world->numTextures, bytes, header->worldSpawn);
    LOAD_CHUNK(ps1bsp_atlas_t, world->atlases, world->numAtlases, bytes, header->atlases);  // numAtlases contains the number of 32-bit words in a single atlas now
    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_polygon_t, world->polygons, world->numPolygons, bytes, header->polygons);
    LOAD_CHUNK(ps1bsp_polyvertex_t, world->polyVertices, world->numPolyVertices, bytes, header->polyVertices);
    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(ps1bsp_model_t, world->models, world->numModels, bytes, header->models);
    LOAD_CHUNK(u_short, world->leafFaces, world->numLeafFaces, bytes, header->leafFaces);
    LOAD_CHUNK(u_char, world->visData, world->numVisData, bytes, header->visData);
}

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_cull_backface(const world_t *world, const ps1bsp_face_t *face)
{
    // Backface culling using the face's plane and center point
    // This eliminates the need for normal clipping per polygon
    VECTOR cam_vec;
    cam_vec.vx = (int)face->center.vx - cam_pos.vx;
    cam_vec.vy = (int)face->center.vy - cam_pos.vy;
    cam_vec.vz = (int)face->center.vz - cam_pos.vz;

    // Check if the face's plane points towards the camera
    const ps1bsp_plane_t *plane = &world->planes[face->planeId];
    int planeDot = m_dot12(&cam_vec, &plane->normal);
    if ((planeDot >= 0) ^ face->side)
        return 0;

    // Check if the face is behind the camera
    // NOTE: disabling the behind-the-camera check does *not* actually solve the problem of polygons disappearing when too close to the camera!
    // This means that the GPU probably already clips polygons that stretch too far outside the drawing area. Tessellation should solve this.
    // UPDATE: tessellation does solve this, to an extent. However we must also make sure not to discard faces partially inside the frustum, and force tessellation on those.
    int camDot = m_dot12(&cam_vec, &cam_dir);
    if (camDot < 0)
    {
        // Check if the face may be partially inside the frustum. If so, we draw it with a high amount of tessellation.
        for (int i = 0; i < 4; ++i)
        {
            const ps1bsp_vertex_t *vert = &world->vertices[face->bounds[i]];
            if (frustum_pointInside((SVECTOR*)vert))
            {
                // TODO: use a simplified formula with only face area to determine tessellation LOD
                return 1;
            }
        }

        // Face really isn't visible
        return 0;
    }

    // Flip angle for faces that are on the opposite side of their plane
    char flip = 2 * face->side - 1;
    planeDot = planeDot * flip;

    // Calculate an LOD value for this face, based on camera position & direction, face normal & area and distance.
    // This is used to determine how many tessellation subdivisions are necessary to prevent texture warping.
    int vecLenSqr = m_dot12(&cam_vec, &cam_vec);
    int lod = (planeDot * vecLenSqr) / (camDot  + 1);
    lod = (lod << 6) / (face->center.pad + 1);
    return (short)lod;
}

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;
}

static void world_drawface_flat(const world_t *world, const ps1bsp_face_t *face, P_COLOR *color, u_long *ot)
{
    // Draw untextured, vertex colored faces, skipping the entire polygon tessellation step
    const ps1bsp_facevertex_t *faceVertices = &world->faceVertices[face->firstFaceVertex];
    draw_quadstrip_flat(world->vertices, faceVertices, face->numFaceVertices, color, ot);
}

static void world_drawface_fast(const world_t *world, const ps1bsp_face_t *face, u_long *ot)
{
    // Draw untextured, vertex colored faces, skipping the entire polygon tessellation step
    const ps1bsp_facevertex_t *faceVertices = &world->faceVertices[face->firstFaceVertex];
    draw_quadstrip_colored(world->vertices, faceVertices, face->numFaceVertices, ot);
}

static void world_drawface_lit(const world_t *world, const ps1bsp_face_t *face, u_long *ot)
{
    // Draw untextured, vertex colored polygons
    const ps1bsp_polygon_t* poly = &world->polygons[face->firstPolygon];
    for (u_char polyIdx = 0; polyIdx < face->numPolygons; ++polyIdx, ++poly)
    {
        const ps1bsp_polyvertex_t *polyVertices = &world->polyVertices[poly->firstPolyVertex];
        draw_quadstrip_lit(world->vertices, polyVertices, poly->numPolyVertices, ot);
    }
}

static void world_drawface_textured(const world_t *world, const ps1bsp_face_t *face, u_long *ot)
{
    // Draw textured, vertex colored polygons
    const ps1bsp_texture_t *texture = &world->textures[face->textureId];
    const ps1bsp_polygon_t* poly = &world->polygons[face->firstPolygon];

    u_char tess = face->flags & 3;
    // if (tess)
    // {
    //     world_drawface_flat(world, face, &colors[tess - 1], ot);
    //     return;
    // }

    if (face->flags & SURF_DRAWLIQUID)
    {
        for (u_char polyIdx = 0; polyIdx < face->numPolygons; ++polyIdx, ++poly)
        {
            const ps1bsp_polyvertex_t *polyVertices = &world->polyVertices[poly->firstPolyVertex];
            draw_quadstrip_water(world->vertices, polyVertices, poly->numPolyVertices, texture->tpage, face->flags & SURF_DRAWWATER, ot);
        }
    }
    else
    {
        for (u_char polyIdx = 0; polyIdx < face->numPolygons; ++polyIdx, ++poly)
        {
            const ps1bsp_polyvertex_t *polyVertices = &world->polyVertices[poly->firstPolyVertex];
            switch (poly->numPolyVertices)
            {
            case 3:
                draw_triangle_textured(world->vertices, polyVertices, texture->tpage, ot);
                break;
            default:
                if (tess)
                    draw_quadstrip_tess2(world->vertices, polyVertices, poly->numPolyVertices, texture->tpage, ot);
                else
                    draw_quadstrip_textured(world->vertices, polyVertices, poly->numPolyVertices, texture->tpage, ot);
                break;
            }
        }
    }

    // Texture window commands needs to be placed *after* the draw commands, because the commands are executed in reverse order
    DR_TWIN *twin = (DR_TWIN*)mem_prim(sizeof(DR_TWIN));
    setlen(twin, 1);
    twin->code[0] = texture->twin;
    addPrim(ot, twin);
}

static void (*world_drawface)(const world_t*, const ps1bsp_face_t*, u_long *ot) = &world_drawface_fast;

static void world_drawFaces(const world_t *world, const ps1bsp_face_t *firstFace)
{
    // Draw the faces in front-to-back order. There are two advantages to this:
    // 1) We don't need to do any depth calculations. The ordering table will be rendered in reverse order, i.e. back-to-front.
    // 2) If we need to stop drawing because the primitive buffer is full, we'll only be culling distant faces.
    for (const ps1bsp_face_t *face = firstFace; face != NULL; face = face->nextFace)
    {
        // Early primitive buffer check
        if (!mem_checkprim(sizeof(POLY_GT4) + sizeof(DR_TWIN), face->totalPrimitives))
            break;

        world_drawface(world, face, curOT);
    }
}

// Simplified BSP subtree traversal specifically for sorting faces of brush models.
// This skips the frustum culling and PVS culling stages, as this will have already been done on the leafs containing the model.
static ps1bsp_face_t *world_sortModelFaces(const world_t* world, const ps1bsp_model_t* model, u_long frameNum, ps1bsp_face_t **lastFace)
{
    short *nodeStack = (short*)(scratchpad + 256);
    u_char stackPtr = 0;

    // Push the start node onto the stack
    nodeStack[stackPtr++] = model->nodeId0;

    ps1bsp_face_t *firstFace = NULL;
    while (stackPtr > 0)
    {
        // Pop a node off the stack
        short nodeIdx = nodeStack[--stackPtr];

        if (nodeIdx < 0)    // Leaf node
        {
            u_short leafIdx = ~nodeIdx;
            if (leafIdx == 0)   // Leaf 0 should not be drawn
                continue;

            const ps1bsp_leaf_t *leaf = &world->leaves[leafIdx];
            const u_short *leafFaces = &world->leafFaces[leaf->firstLeafFace];
            for (u_short leafFaceIdx = 0; leafFaceIdx < leaf->numLeafFaces; ++leafFaceIdx)
            {
                ps1bsp_face_t *face = &world->faces[leafFaces[leafFaceIdx]];

                // Make sure we draw each face only once per frame
                if (face->drawFrame == frameNum)
                    continue;

                face->drawFrame = frameNum;

                // Cull faces that are facing away from the camera
                if (!world_cull_backface(world, face))
                    continue;

                if (!firstFace)
                    *lastFace = face;

                face->nextFace = firstFace;
                firstFace = face;
            }
            continue;
        }

        const ps1bsp_node_t *node = &world->nodes[nodeIdx];
        const ps1bsp_plane_t *plane = &world->planes[node->planeId];
        short dist = world_pointPlaneDist(&cam_pos, plane); // TODO: add/subtract model origin from camera to sort faces in model space
        
        // Traverse the subtree in back-to-front order
        char order = dist < 0;
        nodeStack[stackPtr++] = node->children[order];
        nodeStack[stackPtr++] = node->children[order ^ 1];
    }

    return firstFace;
}

static ps1bsp_face_t *world_sortFaces(const world_t *world, const ps1bsp_leaf_t *firstLeaf)
{
    ps1bsp_face_t *firstFace = NULL;
    u_long frameNum = time_getFrameNumber();

    // Traverse leaves in back-to-front order
    // This ensures that faces shared between adjacent leaves are drawn in the correct order, with the back leaf being drawn behind
    for (const ps1bsp_leaf_t *leaf = firstLeaf; leaf != NULL; leaf = leaf->nextLeaf)
    {
        const u_short *leafFaces = &world->leafFaces[leaf->firstLeafFace];
        for (u_short leafFaceIdx = 0; leafFaceIdx < leaf->numLeafFaces; ++leafFaceIdx)
        {
            ps1bsp_face_t *face = &world->faces[leafFaces[leafFaceIdx]];

            // Make sure we draw each face only once per frame
            if (face->drawFrame == frameNum)
                continue;

            face->drawFrame = frameNum;

            // Cull faces that are facing away from the camera
            // This also returns an LOD value used to determine tessellation level
            short lod = world_cull_backface(world, face);
            if (!lod)
                continue;

            face->flags &= ~3;
            if (lod < 64)
                face->flags |= 1;
            else if (lod < 128)
                face->flags |= 2;
            // else if (lod < 256)
            //     face->flags |= 3;

            // Sort the faces to draw in front-to-back order
            face->nextFace = firstFace;
            firstFace = face;
        }

        // Sort faces for the models in this leaf
        for (const ps1bsp_model_t* model = leaf->models; model != NULL; model = model->nextModel)
        {
            // This successfully sorts the faces within a single model.
            // However, several models occupying the same leaf will still clip through each other
            // and since we're looking at only a single leaf per model, large brush models will often have neighbouring leafs drawing over them.
            ps1bsp_face_t* lastFace;
            ps1bsp_face_t* face = world_sortModelFaces(world, model, frameNum, &lastFace);
            if (face != NULL)
            {
                lastFace->nextFace = firstFace;
                firstFace = face;
            }
        }
    }

    return firstFace;
}

static void world_sortModels(const world_t *world)
{
    for (u_short modelIdx = 1; modelIdx < world->numModels; ++modelIdx)
    {
        ps1bsp_model_t* model = (ps1bsp_model_t*)&world->models[modelIdx];
        ps1bsp_leaf_t* leaf = (ps1bsp_leaf_t*)model->currentLeaf;

        // TODO: skip clip brush-only models (clip and trigger textures)
        // TODO: frustum culling (may be more expensive than it's worth)
        // TODO: instead of leaf, refer to a node, i.e. the first one that splits the model's bounding box/sphere (R_FindTopnode in Q2 code)
        // Q2 distinguishes between models that fall entirely in one leaf (i.e. small ones) and ones split by a node; the latter is clipped by R_DrawSolidClippedSubmodelPolygons
        // Could probably traverse BSP tree from top node and divide up model faces per leaf, insert them as the leaf faces are sorted

        // Update the model's current leaf. This only needs to be done when the model moves.
        if (leaf == NULL)
        {
            VECTOR pos;
            pos.vx = model->origin.vx + model->boundingSphere.vx;
            pos.vy = model->origin.vy + model->boundingSphere.vy;
            pos.vz = model->origin.vz + model->boundingSphere.vz;

            short leafIdx = world_leafAtPoint(world, &pos);
            if (leafIdx <= 0)
                continue;
                
            leaf = (ps1bsp_leaf_t*)&world->leaves[leafIdx];
            model->currentLeaf = leaf;
        }

        model->nextModel = leaf->models;
        leaf->models = model;
    }
}

static ps1bsp_leaf_t *world_sortLeafs(const world_t *world, short startNode, u_char *pvs)
{
    short *nodeStack = (short*)(scratchpad + 256);  // Place the node stack in fast RAM after the PVS data
    u_char stackPtr = 0;

    // Push the start node onto the stack
    nodeStack[stackPtr++] = startNode;

    ps1bsp_leaf_t *firstLeaf = NULL;
    while (stackPtr > 0)
    {
        // Pop a node off the stack
        short nodeIdx = nodeStack[--stackPtr];

        if (nodeIdx < 0)    // Leaf node
        {
            u_short leafIdx = ~nodeIdx;
            if (leafIdx == 0)   // Leaf 0 should not be drawn
                continue;

            // PVS culling
            u_short test = leafIdx - 1;
            if ((pvs[test >> 3] & (1 << (test & 0x7))) == 0)
                continue;

            // Add the leaf to the sorted linked list
            // Since we're traversing the BSP tree front-to-back, adding each leaf at the start sorts the list in back-to-front order
            ps1bsp_leaf_t *leaf = (ps1bsp_leaf_t*)&world->leaves[leafIdx];
            leaf->nextLeaf = firstLeaf;
            leaf->models = NULL;
            firstLeaf = leaf;
            continue;
        }

        // Perform frustum culling
        const ps1bsp_node_t *node = &world->nodes[nodeIdx];
        if (!frustum_sphereInside(&node->boundingSphere))
            continue;

        const ps1bsp_plane_t *plane = &world->planes[node->planeId];
        short dist = world_pointPlaneDist(&cam_pos, plane);
        
        // Traverse the BSP tree in front-to-back order
        char order = dist < 0;
        nodeStack[stackPtr++] = node->children[order ^ 1];
        nodeStack[stackPtr++] = node->children[order];
    }

    return firstLeaf;
}

// 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;
}

void world_draw(const world_t *world)
{
    // 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 + 48;   // Place the PVS data after the frustum in fast RAM
    u_char *pvs = world_loadVisData(world, cam_leaf, &pvsbuf);
    //u_char *pvs = world_noVisData(world, &pvsbuf);

    ps1bsp_leaf_t *firstLeaf = world_sortLeafs(world, world->models[0].nodeId0, pvs);
    world_sortModels(world);
    ps1bsp_face_t *firstFace = world_sortFaces(world, firstLeaf);

    if (enableTexturing)
        world_drawface = &world_drawface_textured;
    else
        world_drawface = &world_drawface_lit;

    // Make sure we set the texture window back to default after drawing the world
    DR_TWIN *twin = (DR_TWIN*)mem_prim(sizeof(DR_TWIN));
    setlen(twin, 1);
    twin->code[0] = 0xe2000000;
    addPrim(curOT, twin);

    world_drawFaces(world, firstFace);
}