#include "crBatchedPrimitives.h"

#include <algorithm>
#include <cstddef>

#include <SDL3/SDL_stdinc.h>

#include "crGraphicsStats.h"
#include "crMath.h"
#include "crPrimitiveGeometry.h"

// view context for the alpha depth comparator — std::sort takes a plain function pointer
// (no lambda captures by project style), so the per-sort camera state lives here
static float3 s_alphaSortEye;
static float3 s_alphaSortFwd;

void crBatchedPrimitives::Init()
{
    {
        crArray<float>    verts;
        crArray<uint16_t> indices;

        crPrimitiveGeometry::BuildBox( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::BOX ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildSphere( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::SPHERE ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildSphereShell( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::SPHERE_SHELL ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildOctahedron( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::OCTAHEDRON ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildIcosahedron( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::ICOSAHEDRON ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildBeam( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::BEAM ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildRing( &verts, &indices, crPrimitiveGeometry::RING_INNER );
        _buffers[ static_cast<int32_t>( EPrimitive::RING ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildRing( &verts, &indices, crPrimitiveGeometry::RING_INNER_THIN );
        _buffers[ static_cast<int32_t>( EPrimitive::RING_THIN ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildTetrahedron( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::TETRAHEDRON ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildArrow( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::ARROW ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildHexagon( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::HEXAGON ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildStar( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::STAR ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildSpike( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::SPIKE ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );

        verts.Clear();
        indices.Clear();
        crPrimitiveGeometry::BuildDiamond( &verts, &indices );
        _buffers[ static_cast<int32_t>( EPrimitive::DIAMOND ) ] = CreateBuffers( verts.Data(), verts.Size(), indices.Data(), indices.Size() );
    }

    const int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );
    for( int32_t i = 0; i < len; ++i )
        SDL_assert( _buffers[ i ].indexCount > 0 );   // an EPrimitive with no generator draws nothing and says nothing
}
uint8_t crBatchedPrimitives::RegisterMaterial( GLuint program )
{
    if( _materialCount >= MAX_MATERIALS )
    {
        SDL_LogError( SDL_LOG_CATEGORY_RENDER, "crBatchedPrimitives: material registry full (%d) — raise MAX_MATERIALS", MAX_MATERIALS );
        SDL_assert( false );
        return 0;
    }

    const uint8_t id = static_cast<uint8_t>( _materialCount );
    _materialProgram[ id ] = program;
    ++_materialCount;
    return id;
}
void crBatchedPrimitives::Cleanup()
{
    int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );
    for( int32_t i = 0; i < len; ++i )
    {
        PrimitiveBuffers& b = _buffers[ i ];
        glDeleteBuffers( 1, &b.instanceVbo );
        glDeleteBuffers( 1, &b.ibo );
        glDeleteBuffers( 1, &b.vbo );
        glDeleteVertexArrays( 1, &b.vao );
        b = {};
    }
}

void crBatchedPrimitives::Begin( const crFrustum& cameraFrustum, const crFrustum* cascadeFrustums )
{
    _cameraFrustum = cameraFrustum;
    for( int32_t c = 0; c < crShadowMap::CASCADE_COUNT; ++c )
        _cascadeFrustums[ c ] = cascadeFrustums[ c ];

    int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );
    for( int32_t i = 0; i < len; ++i )
    {
        for( int32_t m = 0; m < _materialCount; ++m )
        {
            _opaque[ m ][ i ].Clear();
            _multiply[ m ][ i ].Clear();
            _additive[ m ][ i ].Clear();
        }
        for( int32_t c = 0; c < crShadowMap::CASCADE_COUNT; ++c )
            _shadow[ c ][ i ].Clear();

    }

    _alpha.Clear();
}
void crBatchedPrimitives::Add( EPrimitive shape, EBlendMode blend, quat4 rot, float3 pos, float3 scale, color4 color, float emissive, float receiveShadow, bool castShadow, float fresnel, float custom0, float custom1, float custom2, float custom3, float custom4, float custom5, uint8_t material )
{
    PrimitiveInstance inst;
    inst.rot           = rot;
    inst.pos           = pos;
    inst.scale         = scale;
    inst.color         = color;
    inst.emissive      = emissive;
    inst.receiveShadow = receiveShadow;
    inst.fresnel       = fresnel;
    inst.custom0       = custom0;
    inst.custom1       = custom1;
    inst.custom2       = custom2;
    inst.custom3       = custom3;
    inst.custom4       = custom4;
    inst.custom5       = custom5;

    const int32_t s      = static_cast<int32_t>( shape );
    const float   radius = crMath::Sqrt( crMath::Dot3( scale, scale ) );   // bounding sphere = box diagonal; conservative for the sphere primitive too, and rotation-invariant so rot is irrelevant

    int32_t m = static_cast<int32_t>( material );
    if( m >= _materialCount )
    {
        SDL_assert( false );   // an unregistered id draws with the wrong shader and says nothing
        m = 0;
    }

    const bool inCamera = crMath::SphereInFrustum( _cameraFrustum, pos, radius );

    // transparent modes never enter the shadow pass — a binary depth map cannot represent them
    switch( blend )
    {
    case EBlendMode::MULTIPLY:
        if( inCamera )
            _multiply[ m ][ s ].Add( inst );
        return;
    case EBlendMode::ALPHA:
        if( inCamera )
        {// sorted-alpha is builtin-only — a custom material can't join the global depth sort
            AlphaEntry entry;
            entry.inst  = inst;
            entry.shape = static_cast<uint8_t>( shape );
            _alpha.Add( entry );
        }
        return;
    case EBlendMode::ADDITIVE:
        if( inCamera )
            _additive[ m ][ s ].Add( inst );
        return;
    default:
        break;
    }

    if( inCamera )
        _opaque[ m ][ s ].Add( inst );

    // shadow casters cull against the LIGHT volumes, not the camera — an off-screen caster can
    // still drop a shadow into view. each cascade keeps only what its own volume covers
    if( castShadow )
    {
        for( int32_t c = 0; c < crShadowMap::CASCADE_COUNT; ++c )
        {
            if( crMath::SphereInFrustum( _cascadeFrustums[ c ], pos, radius ) )
                _shadow[ c ][ s ].Add( inst );
        }
    }
}
void crBatchedPrimitives::Flush( GLuint builtinProgram )
{
    _materialProgram[ 0 ] = builtinProgram;

    glEnable( GL_DEPTH_TEST );
    glDisable( GL_BLEND );      // opaque pass
    glEnable( GL_CULL_FACE );   // CCW front faces — walls are one-sided (invisible once the camera is inside)
    glCullFace( GL_BACK );      // declare, don't inherit — cull mode is sticky GL state
    glEnable( GL_POLYGON_OFFSET_FILL );
    glPolygonOffset( 1.0f, 1.0f );   // push surfaces back so coplanar debug wireframes win the depth test

    const int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );

    for( int32_t m = 0; m < _materialCount; ++m )
    {
        if( _materialProgram[ m ] == 0 )
            continue;
        glUseProgram( _materialProgram[ m ] );
        for( int32_t i = 0; i < len; ++i )
            DrawInstances( _buffers[ i ], _opaque[ m ][ i ] );
    }

    glBindVertexArray( 0 );

    glDisable( GL_POLYGON_OFFSET_FILL );
    glDisable( GL_CULL_FACE );
}
void crBatchedPrimitives::FlushShadow( GLuint program, int32_t cascade )
{
    glUseProgram( program );

    glEnable( GL_DEPTH_TEST );
    glDisable( GL_BLEND );
    glEnable( GL_CULL_FACE );
    glCullFace( GL_BACK );    // record FRONT faces — back-face storage was tried twice: on this contact-heavy
                              // content it erodes contact shadows and leaks light at touching surfaces
    glEnable( GL_POLYGON_OFFSET_FILL );
    glPolygonOffset( 2.0f, 4.0f );   // slope-scaled: grazing faces need a depth margin proportional to their slope —
                                     // a constant bias sized for the worst angle shrank shadow coverage everywhere

    int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );
    for( int32_t i = 0; i < len; ++i )
        DrawInstances( _buffers[ i ], _shadow[ cascade ][ i ] );

    glBindVertexArray( 0 );

    glDisable( GL_POLYGON_OFFSET_FILL );
    glDisable( GL_CULL_FACE );
}
void crBatchedPrimitives::FlushTransparent( GLuint builtinProgram, float3 eye, float3 fwd )
{
    _materialProgram[ 0 ] = builtinProgram;

    glEnable( GL_DEPTH_TEST );
    glDepthMask( GL_FALSE );    // read depth (occluded by opaque) but never write — transparents must not erase each other
    glEnable( GL_BLEND );
    glEnable( GL_CULL_FACE );
    glCullFace( GL_BACK );

    int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );

    // MULTIPLY: darkening — multiplication commutes, no sorting needed
    glBlendFunc( GL_DST_COLOR, GL_ZERO );
    for( int32_t m = 0; m < _materialCount; ++m )
    {
        if( _materialProgram[ m ] == 0 )
            continue;
        glUseProgram( _materialProgram[ m ] );
        for( int32_t i = 0; i < len; ++i )
            DrawInstances( _buffers[ i ], _multiply[ m ][ i ] );
    }

    glUseProgram( builtinProgram );   // ALPHA + the sort below are builtin-only

    // ALPHA: the only order-dependent mode — one GLOBAL back-to-front sort across every primitive,
    // drawn as instanced runs (one draw per consecutive same-primitive stretch). correctness never
    // degrades; only the draw count grows when primitives interleave in depth
    glBlendFunc( GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA );

    const int32_t alphaCount = _alpha.Size();
    if( alphaCount > 0 )
    {
        s_alphaSortEye = eye;
        s_alphaSortFwd = fwd;
        if( alphaCount > 1 )
            std::sort( _alpha.Data(), _alpha.Data() + alphaCount, &crBatchedPrimitives::AlphaDepthGreater );

        // fill the per-primitive VBOs in global order — every run then maps to a consecutive VBO range
        for( int32_t i = 0; i < len; ++i )
            _alphaSorted[ i ].Clear();
        for( int32_t i = 0; i < alphaCount; ++i )
            _alphaSorted[ _alpha.At( i ).shape ].Add( _alpha.At( i ).inst );

        for( int32_t i = 0; i < len; ++i )
        {
            const int32_t count = _alphaSorted[ i ].Size();
            if( count <= 0 )
                continue;

            glBindVertexArray( _buffers[ i ].vao );
            glBindBuffer( GL_ARRAY_BUFFER, _buffers[ i ].instanceVbo );
            glBufferData( GL_ARRAY_BUFFER, static_cast<GLsizeiptr>( count ) * static_cast<GLsizeiptr>( sizeof( PrimitiveInstance ) ), _alphaSorted[ i ].Data(), GL_STREAM_DRAW );
        }

        int32_t cursor[ static_cast<int32_t>( EPrimitive::_SIZE ) ] = {};

        int32_t i = 0;
        while( i < alphaCount )
        {
            const uint8_t shape = _alpha.At( i ).shape;

            int32_t runEnd = i + 1;
            while( ( runEnd < alphaCount ) && ( _alpha.At( runEnd ).shape == shape ) )
                ++runEnd;

            const int32_t runLength = runEnd - i;
            DrawInstanceRange( _buffers[ shape ], cursor[ shape ], runLength );
            cursor[ shape ] += runLength;
            i = runEnd;
        }
    }

    // ADDITIVE: glow — addition commutes; HDR accumulates past 1.0 straight into bloom
    glBlendFunc( GL_ONE, GL_ONE );
    for( int32_t m = 0; m < _materialCount; ++m )
    {
        if( _materialProgram[ m ] == 0 )
            continue;
        glUseProgram( _materialProgram[ m ] );
        for( int32_t i = 0; i < len; ++i )
            DrawInstances( _buffers[ i ], _additive[ m ][ i ] );
    }

    glBindVertexArray( 0 );

    glDepthMask( GL_TRUE );   // sticky state — the shadow/scene depth clears next frame need writes back on
    glDisable( GL_CULL_FACE );
}
void crBatchedPrimitives::ClearUnusedMemory()
{
    const int32_t len = static_cast<int32_t>( EPrimitive::_SIZE );
    for( int32_t i = 0; i < len; ++i )
    {
        for( int32_t m = 0; m < _materialCount; ++m )
        {
            TrimInstances( &_opaque[ m ][ i ] );
            TrimInstances( &_multiply[ m ][ i ] );
            TrimInstances( &_additive[ m ][ i ] );
        }

        for( int32_t c = 0; c < crShadowMap::CASCADE_COUNT; ++c )
            TrimInstances( &_shadow[ c ][ i ] );

        TrimInstances( &_alphaSorted[ i ] );
    }

    if( _alpha.Capacity() > 0 )
    {
        _alpha.Clear();
        _alpha.Shrink();
    }
}
void crBatchedPrimitives::DrawInstances( const PrimitiveBuffers& buffers, const crArray<PrimitiveInstance>& instances )
{
    const int32_t count = instances.Size();
    if( count <= 0 )
        return;

    glBindVertexArray( buffers.vao );
    glBindBuffer( GL_ARRAY_BUFFER, buffers.instanceVbo );
    glBufferData( GL_ARRAY_BUFFER, static_cast<GLsizeiptr>( count ) * static_cast<GLsizeiptr>( sizeof( PrimitiveInstance ) ), instances.Data(), GL_STREAM_DRAW );

    DrawInstanceRange( buffers, 0, count );
}
void crBatchedPrimitives::DrawInstanceRange( const PrimitiveBuffers& buffers, int32_t first, int32_t count )
{
    glBindVertexArray( buffers.vao );
    glBindBuffer( GL_ARRAY_BUFFER, buffers.instanceVbo );
    SetInstanceAttribPointers( first );   // every draw declares its own base — alpha runs leave the VAO mid-buffer

    glDrawElementsInstanced( GL_TRIANGLES, buffers.indexCount, GL_UNSIGNED_SHORT, nullptr, count );

    ++crGraphicsStats::drawCalls;
    crGraphicsStats::vertices += ( buffers.indexCount * count );
}
/*static*/ void crBatchedPrimitives::SetInstanceAttribPointers( int32_t firstInstance )
{
    // PrimitiveInstance is a POD laid out to feed vertex attributes directly; offsetof keeps the pointer
    // offsets bound to the struct so a field reorder can't silently desync them. the static_assert
    // pins the total size — if it trips, the layout changed and the shaders/attribs need review
    static_assert( sizeof( PrimitiveInstance ) == 92, "PrimitiveInstance layout changed — resync attribs + primitive shaders" );

    const GLsizei   stride = static_cast<GLsizei>( sizeof( PrimitiveInstance ) );
    const uintptr_t base   = static_cast<uintptr_t>( firstInstance ) * sizeof( PrimitiveInstance );

    glVertexAttribPointer( 2, 4, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, rot ) ) );
    glVertexAttribPointer( 3, 3, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, pos ) ) );
    glVertexAttribPointer( 4, 3, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, scale ) ) );
    glVertexAttribPointer( 5, 4, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, color ) ) );
    glVertexAttribPointer( 6, 4, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, emissive ) ) );   // emissive, receiveShadow, fresnel, custom0
    glVertexAttribPointer( 7, 4, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, custom1 ) ) );    // custom1..4
    glVertexAttribPointer( 8, 1, GL_FLOAT, GL_FALSE, stride, reinterpret_cast<void*>( base + offsetof( PrimitiveInstance, custom5 ) ) );    // custom5
}

/*static*/ bool crBatchedPrimitives::AlphaDepthGreater( const AlphaEntry& a, const AlphaEntry& b )
{
    const float da = ( ( a.inst.pos.x - s_alphaSortEye.x ) * s_alphaSortFwd.x )
                   + ( ( a.inst.pos.y - s_alphaSortEye.y ) * s_alphaSortFwd.y )
                   + ( ( a.inst.pos.z - s_alphaSortEye.z ) * s_alphaSortFwd.z );
    const float db = ( ( b.inst.pos.x - s_alphaSortEye.x ) * s_alphaSortFwd.x )
                   + ( ( b.inst.pos.y - s_alphaSortEye.y ) * s_alphaSortFwd.y )
                   + ( ( b.inst.pos.z - s_alphaSortEye.z ) * s_alphaSortFwd.z );
    return da > db;
}
/*static*/ void crBatchedPrimitives::TrimInstances( crArray<PrimitiveInstance>* instances )
{
    instances->Clear();
    instances->Shrink();
}
crBatchedPrimitives::PrimitiveBuffers crBatchedPrimitives::CreateBuffers( const float* verts, int32_t floatCount, const uint16_t* indices, int32_t indexCount )
{
    PrimitiveBuffers b = {};
    b.indexCount = indexCount;

    glGenVertexArrays( 1, &b.vao );
    glGenBuffers( 1, &b.vbo );
    glGenBuffers( 1, &b.ibo );
    glGenBuffers( 1, &b.instanceVbo );

    glBindVertexArray( b.vao );

    glBindBuffer( GL_ARRAY_BUFFER, b.vbo );
    glBufferData( GL_ARRAY_BUFFER, static_cast<GLsizeiptr>( floatCount ) * static_cast<GLsizeiptr>( sizeof( float ) ), verts, GL_STATIC_DRAW );

    const GLsizei vertStride = static_cast<GLsizei>( 6 * sizeof( float ) );
    glEnableVertexAttribArray( 0 );
    glVertexAttribPointer( 0, 3, GL_FLOAT, GL_FALSE, vertStride, reinterpret_cast<void*>( 0 ) );
    glEnableVertexAttribArray( 1 );
    glVertexAttribPointer( 1, 3, GL_FLOAT, GL_FALSE, vertStride, reinterpret_cast<void*>( 12 ) );

    glBindBuffer( GL_ELEMENT_ARRAY_BUFFER, b.ibo );
    glBufferData( GL_ELEMENT_ARRAY_BUFFER, static_cast<GLsizeiptr>( indexCount ) * static_cast<GLsizeiptr>( sizeof( uint16_t ) ), indices, GL_STATIC_DRAW );

    glBindBuffer( GL_ARRAY_BUFFER, b.instanceVbo );

    glEnableVertexAttribArray( 2 );
    glEnableVertexAttribArray( 3 );
    glEnableVertexAttribArray( 4 );
    glEnableVertexAttribArray( 5 );
    glEnableVertexAttribArray( 6 );
    glEnableVertexAttribArray( 7 );
    glEnableVertexAttribArray( 8 );
    SetInstanceAttribPointers( 0 );
    glVertexAttribDivisor( 2, 1 );
    glVertexAttribDivisor( 3, 1 );
    glVertexAttribDivisor( 4, 1 );
    glVertexAttribDivisor( 5, 1 );
    glVertexAttribDivisor( 6, 1 );
    glVertexAttribDivisor( 7, 1 );
    glVertexAttribDivisor( 8, 1 );

    glBindVertexArray( 0 );
    glBindBuffer( GL_ARRAY_BUFFER, 0 );
    glBindBuffer( GL_ELEMENT_ARRAY_BUFFER, 0 );   // after the VAO unbind — the ELEMENT binding is VAO state

    return b;
}