#include <iostream>
#include <vector>
#include <fstream>
#include <limits>
#include <stdint.h>

#include <assimp/Importer.hpp>      // C++ importer interface
#include <assimp/scene.h>           // Output data structure
#include <assimp/postprocess.h>     // Post processing flags
#include <assimp/ai_assert.h>

#pragma comment (lib, "assimp.lib")

// One instance per aiMesh in the globally loaded asset
MeshHelper** apcMeshes;

const aiScene* scene;

////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////
bool Load(const std::string& pFile)
{
    // Create an instance of the Importer class
    Assimp::Importer importer;
    // And have it read the given file with some example postprocessing
    // Usually - if speed is not the most important aspect for you - you'll
    // propably to request more postprocessing than we do in this example.
    scene = importer.ReadFile(pFile,
                            aiProcess_CalcTangentSpace       |
                            aiProcess_Triangulate            |
                            aiProcess_JoinIdenticalVertices  |
                            aiProcess_FlipUVs);

    // If the import failed, report it
    if(!scene)
    {
        std::string s = "Error in load the " + pFile;
        MessageBoxA(NULL, s.c_str(), "Error mesh", MB_OK);
        return false;
    }

    // allocate a new MeshHelper array and build a new instance
    // for each mesh in the original asset
    apcMeshes = new MeshHelper*[scene->mNumMeshes]();
    for (unsigned int i = 0; i < scene->mNumMeshes;++i)
        apcMeshes[i] = new MeshHelper();

    for (unsigned int i = 0; i < scene->mNumMeshes;++i)
    {
        const aiMesh* mesh = scene->mMeshes[i];

        /*
        // create the material for the mesh
        if (!apcMeshes[i]->piEffect)    {
            CMaterialManager::Instance().CreateMaterial(apcMeshes[i], mesh);
        }
        */

        // create vertex buffer
        if(FAILED( devices.d3ddev->CreateVertexBuffer(sizeof(Vertex) *
                                                        mesh->mNumVertices,
                                                        D3DUSAGE_WRITEONLY,
                                                        0,
                                                        D3DPOOL_DEFAULT, &apcMeshes[i]->piVB,NULL)))
        {
            MessageBoxA(NULL, "Failed to create vertex buffer", "Error Model", MB_OK);
            return false;
        }

        DWORD dwUsage = 0;
        if (apcMeshes[i]->piOpacityTexture || 1.0f != apcMeshes[i]->fOpacity)
            dwUsage |= D3DUSAGE_DYNAMIC;

        unsigned int nidx;
        switch (mesh->mPrimitiveTypes)
        {
            case aiPrimitiveType_POINT:
                nidx = 1;
                break;
            case aiPrimitiveType_LINE:
                nidx = 2;
                break;
            case aiPrimitiveType_TRIANGLE:
                nidx = 3;
                break;
            default:
                ai_assert(false);
        };


        // check whether we can use 16 bit indices
        if (mesh->mNumFaces * 3 >= 65536)   {
            // create 32 bit index buffer
            if(FAILED( devices.d3ddev->CreateIndexBuffer( 4 * mesh->mNumFaces * nidx,
                                                        D3DUSAGE_WRITEONLY | dwUsage,
                                                        D3DFMT_INDEX32,
                                                        D3DPOOL_DEFAULT,
                                                        &apcMeshes[i]->piIB,
                                                        NULL)))
            {
                MessageBoxA(NULL, "Failed to create 32 Bit index buffer", "Error model", MB_OK);
                return false;
            }

            // now fill the index buffer
            unsigned int* pbData;
            apcMeshes[i]->piIB->Lock(0,0,(void**)&pbData,0);
            for (unsigned int x = 0; x < mesh->mNumFaces;++x)
            {
                for (unsigned int a = 0; a < nidx;++a)
                {
                    *pbData++ = mesh->mFaces[x].mIndices[a];
                }
            }
        }
        else    {
            // create 16 bit index buffer
            if(FAILED( devices.d3ddev->CreateIndexBuffer( 2 * mesh->mNumFaces * nidx,
                                                        D3DUSAGE_WRITEONLY | dwUsage,
                                                        D3DFMT_INDEX16,
                                                        D3DPOOL_DEFAULT,
                                                        &apcMeshes[i]->piIB,
                                                        NULL)))
            {
                MessageBoxA(NULL, "Failed to create 16 Bit index buffer", "Error model", MB_OK);
                return false;
            }

            // now fill the index buffer
            uint16_t* pbData;
            apcMeshes[i]->piIB->Lock(0,0,(void**)&pbData,0);
            for (unsigned int x = 0; x < mesh->mNumFaces;++x)
            {
                for (unsigned int a = 0; a < nidx;++a)
                {
                    *pbData++ = (uint16_t)mesh->mFaces[x].mIndices[a];
                }
            }
        }
        apcMeshes[i]->piIB->Unlock();

        // collect weights on all vertices. Quick and careless
        std::vector<std::vector<aiVertexWeight>> weightsPerVertex( mesh->mNumVertices);
        for( unsigned int a = 0; a < mesh->mNumBones; a++)  {
            const aiBone* bone = mesh->mBones[a];
            for( unsigned int b = 0; b < bone->mNumWeights; b++)
                weightsPerVertex[bone->mWeights[b].mVertexId].push_back( aiVertexWeight( a, bone->mWeights[b].mWeight));
        }

        // now fill the vertex buffer
        Vertex* pbData2;
        apcMeshes[i]->piVB->Lock(0,0,(void**)&pbData2,0);
        for (unsigned int x = 0; x < mesh->mNumVertices;++x)
        {
            pbData2->vPosition = mesh->mVertices[x];

            if (NULL == mesh->mNormals)
                pbData2->vNormal = aiVector3D(0.0f,0.0f,0.0f);
            else pbData2->vNormal = mesh->mNormals[x];

            if (NULL == mesh->mTangents)    {
                pbData2->vTangent = aiVector3D(0.0f,0.0f,0.0f);
                pbData2->vBitangent = aiVector3D(0.0f,0.0f,0.0f);
            }
            else    {
                pbData2->vTangent = mesh->mTangents[x];
                pbData2->vBitangent = mesh->mBitangents[x];
            }

            if (mesh->HasVertexColors( 0 )) {
                using std::min;
                using std::max;
                unsigned char m_argb[4];
                m_argb[0] = (unsigned char)max( min( mesh->mColors[0][x].a * 255.0f, 255.0f), 0.0f);
                m_argb[1] = (unsigned char)max( min( mesh->mColors[0][x].r * 255.0f, 255.0f), 0.0f);
                m_argb[2] = (unsigned char)max( min( mesh->mColors[0][x].g * 255.0f, 255.0f), 0.0f);
                m_argb[3] = (unsigned char)max( min( mesh->mColors[0][x].b * 255.0f, 255.0f), 0.0f);

                pbData2->dColorDiffuse = D3DCOLOR_ARGB(m_argb[0],m_argb[1],m_argb[2],m_argb[3]);
            }
            else pbData2->dColorDiffuse = D3DCOLOR_ARGB(0xFF,0xff,0xff,0xff);

            // ignore a third texture coordinate component
            if (mesh->HasTextureCoords( 0)) {
                pbData2->vTextureUV = aiVector2D(
                    mesh->mTextureCoords[0][x].x,
                    mesh->mTextureCoords[0][x].y);
            }
            else pbData2->vTextureUV = aiVector2D(0.5f,0.5f);

            if (mesh->HasTextureCoords( 1)) {
                pbData2->vTextureUV2 = aiVector2D(
                    mesh->mTextureCoords[1][x].x,
                    mesh->mTextureCoords[1][x].y);
            }
            else pbData2->vTextureUV2 = aiVector2D(0.5f,0.5f);

            // Bone indices and weights
            if( mesh->HasBones())   {
                unsigned char boneIndices[4] = { 0, 0, 0, 0 };
                unsigned char boneWeights[4] = { 0, 0, 0, 0 };
                ai_assert( weightsPerVertex[x].size() <= 4);
                for( unsigned int a = 0; a < weightsPerVertex[x].size(); a++)
                {
                    boneIndices[a] = weightsPerVertex[x][a].mVertexId;
                    boneWeights[a] = (unsigned char) (weightsPerVertex[x][a].mWeight * 255.0f);
                }

                memcpy( pbData2->mBoneIndices, boneIndices, sizeof( boneIndices));
                memcpy( pbData2->mBoneWeights, boneWeights, sizeof( boneWeights));
            } else
            {
                memset( pbData2->mBoneIndices, 0, sizeof( pbData2->mBoneIndices));
                memset( pbData2->mBoneWeights, 0, sizeof( pbData2->mBoneWeights));
            }

            ++pbData2;
        }
        apcMeshes[i]->piVB->Unlock();

        // now generate the second vertex buffer, holding all normals
        if (!apcMeshes[i]->piVBNormals) {
            //GenerateNormalsAsLineList(apcMeshes[i],mesh);

            ai_assert(NULL != apcMeshes[i]);
            ai_assert(NULL != mesh);

            if (!mesh->mNormals)return 0;

            // create vertex buffer
            if(FAILED( devices.d3ddev->CreateVertexBuffer(sizeof(LineVertex) *
                                                        mesh->mNumVertices * 2,
                                                        D3DUSAGE_WRITEONLY,
                                                        LineVertex::GetFVF(),
                                                        D3DPOOL_DEFAULT, &apcMeshes[i]->piVBNormals,NULL)))
            {
                MessageBoxA(NULL,"Failed to create vertex buffer for the normal list", "Error model", MB_OK);
                return false;
            }

            // now fill the vertex buffer with data
            LineVertex* pbData2;
            apcMeshes[i]->piVBNormals->Lock(0,0,(void**)&pbData2,0);
            for (unsigned int x = 0; x < mesh->mNumVertices;++x)
            {
                pbData2->vPosition = mesh->mVertices[x];

                ++pbData2;

                aiVector3D vNormal = mesh->mNormals[x];
                vNormal.Normalize();

                // scalo with the inverse of the world scaling to make sure
                // the normals have equal length in each case
                // TODO: Check whether this works in every case, I don't think so
                //vNormal.x /= g_mWorld.a1*4;
                //vNormal.y /= g_mWorld.b2*4;
                //vNormal.z /= g_mWorld.c3*4;

                pbData2->vPosition = mesh->mVertices[x] + vNormal;

                ++pbData2;
            }
            apcMeshes[i]->piVBNormals->Unlock();
        }
    }

    // get the number of vertices/faces in the model
    unsigned int iNumVert = 0;
    unsigned int iNumFaces = 0;
    for (unsigned int i = 0; i < scene->mNumMeshes;++i)
    {
        iNumVert += scene->mMeshes[i]->mNumVertices;
        iNumFaces += scene->mMeshes[i]->mNumFaces;
    }

    meshrender = new MeshRender(devices, scene, apcMeshes);

    return true;
}

////////////////////////////////////////////////////////////
// Render
////////////////////////////////////////////////////////////
void Render(bool wireframe)
{
    if(wireframe)
    {
        devices.d3ddev->SetRenderState(D3DRS_FILLMODE, D3DFILL_WIREFRAME);
    }

    aiMatrix4x4 m;
    RenderNode(scene->mRootNode, true);
    // render all child nodes
    for (unsigned int i = 0; i < scene->mRootNode->mNumChildren;++i)
        RenderNode(scene->mRootNode->mChildren[i], true );

    if(wireframe)
    {
        devices.d3ddev->SetRenderState(D3DRS_FILLMODE, D3DFILL_SOLID);
    }
}

////////////////////////////////////////////////////////////
// Render a single node
////////////////////////////////////////////////////////////
bool RenderNode(aiNode* piNode, bool bAlpha)
{
    aiMatrix4x4 aiMe;
    aiMatrix4x4* piMatrix = new aiMatrix4x4();

    for (unsigned int i = 0; i < piNode->mNumMeshes;++i)
    {
        if (bAlpha)
            meshrender->DrawSorted(piNode->mMeshes[i], aiMe);
        else
            meshrender->DrawUnsorted(piNode->mMeshes[i]);
    }


    return true;
}

////////////////////////////////////////////////////////////
// Release
////////////////////////////////////////////////////////////
void Release()
{
    bool bNoMaterials = false;
    if (!scene)
        return;

    // TODO: Move this to a proper destructor
    for (unsigned int i = 0; i < scene->mNumMeshes;++i)
    {
        if(apcMeshes[i]->piVB)
        {
            apcMeshes[i]->piVB->Release();
            apcMeshes[i]->piVB = NULL;
        }
        if(apcMeshes[i]->piVBNormals)
        {
            apcMeshes[i]->piVBNormals->Release();
            apcMeshes[i]->piVBNormals = NULL;
        }
        if(apcMeshes[i]->piIB)
        {
            apcMeshes[i]->piIB->Release();
            apcMeshes[i]->piIB = NULL;
        }

        if (!bNoMaterials)
        {
            if(apcMeshes[i]->piEffect)
            {
                apcMeshes[i]->piEffect->Release();
                apcMeshes[i]->piEffect = NULL;
            }
            if(apcMeshes[i]->piDiffuseTexture)
            {
                apcMeshes[i]->piDiffuseTexture->Release();
                apcMeshes[i]->piDiffuseTexture = NULL;
            }
            if(apcMeshes[i]->piNormalTexture)
            {
                apcMeshes[i]->piNormalTexture->Release();
                apcMeshes[i]->piNormalTexture = NULL;
            }
            if(apcMeshes[i]->piSpecularTexture)
            {
                apcMeshes[i]->piSpecularTexture->Release();
                apcMeshes[i]->piSpecularTexture = NULL;
            }
            if(apcMeshes[i]->piAmbientTexture)
            {
                apcMeshes[i]->piAmbientTexture->Release();
                apcMeshes[i]->piAmbientTexture = NULL;
            }
            if(apcMeshes[i]->piEmissiveTexture)
            {
                apcMeshes[i]->piEmissiveTexture->Release();
                apcMeshes[i]->piEmissiveTexture = NULL;
            }
            if(apcMeshes[i]->piOpacityTexture)
            {
                apcMeshes[i]->piOpacityTexture->Release();
                apcMeshes[i]->piOpacityTexture = NULL;
            }
            if(apcMeshes[i]->piShininessTexture)
            {
                apcMeshes[i]->piShininessTexture->Release();
                apcMeshes[i]->piShininessTexture = NULL;
            }
        }
    }
}