#include "StaticMesh.h"
#include "ObjImportUtil.h"

StaticMesh::StaticMesh()
    : mesh(nullptr)
{
    memset(shaderData, 0
        , sizeof(DKObject<DKData>) * static_cast<int>(DKShaderStage::Compute));
}

bool StaticMesh::LoadMeshResourceFromFile(DKObject<DKGraphicsDevice> inDevice
    , DKResourcePool& inDKResourcePool, const DKString& inFileName)
{
    if (inDevice == nullptr)
        return false;

    if (inFileName.Length() == 0)
        return false;

    if (inFileName.HasSuffix("obj"))
    {
        ObjImportUtil::ObjImportedMeshData loadedData
            = ObjImportUtil::LoadFromObjFile(inFileName, inDKResourcePool);

        mesh = DKObject<DKMesh>::New();

        uint32_t vertexBufferSize
            = static_cast<uint32_t>(loadedData.vertices.Count()) * loadedData.vertexSize;

        DKVertexBuffer vb = {};
        {
            DKObject<DKGpuBuffer> vertexBuffer
                = inDevice->CreateBuffer(vertexBufferSize, DKGpuBuffer::StorageModeShared, DKCpuCacheModeReadWrite);
            memcpy(vertexBuffer->Contents(), loadedData.vertices, vertexBufferSize);
            vertexBuffer->Flush();

            vb.buffer = vertexBuffer;
            vb.count = loadedData.vertices.Count();
            vb.offset = 0;
            vb.size = loadedData.vertexSize;
            vb.declarations =
            {
                {DKVertexStream::Position, DKVertexFormat::Float3, false, offsetof(ObjImportUtil::ObjVertex, inPos), DKString("Position")},
                {DKVertexStream::Color, DKVertexFormat::Float3, false, offsetof(ObjImportUtil::ObjVertex, inColor), DKString("Color")},
                {DKVertexStream::Normal, DKVertexFormat::Float3, false, offsetof(ObjImportUtil::ObjVertex, inNormal), DKString("Normal")},
                {DKVertexStream::TexCoord, DKVertexFormat::Float2, false, offsetof(ObjImportUtil::ObjVertex, inTexCoord), DKString("TexCoord")},
            };

            mesh->vertexBuffers.Add(vb);
        }

        DKSubMesh subMesh = {};
        {
            uint32_t indexBufferSize
                = loadedData.indices.Count() * sizeof(uint32_t);

            DKObject<DKGpuBuffer> indexBuffer =
                inDevice->CreateBuffer(indexBufferSize
                    , DKGpuBuffer::StorageModeShared, DKCpuCacheModeReadWrite);

            memcpy(indexBuffer->Contents(), loadedData.indices, indexBufferSize);
            indexBuffer->Flush();

            subMesh.indexBuffer = indexBuffer;
            subMesh.indexCount = loadedData.indices.Count();
            subMesh.indexOffset = 0;
            subMesh.indexType = DKIndexType::UInt32;
            subMesh.vertexOffset = 0;
            subMesh.visible = true;
        }

        mesh->subMeshes.Add(subMesh);

        aabb = loadedData.aabb;

        return true;
    }

    return false;
}

bool StaticMesh::LoadTextureFromFile(DKResourcePool& inDKResourcePool
    , const DKString& inFileName)
{
   textureSrc = DKImage::Create(inDKResourcePool.LoadResourceData(inFileName));
   if (textureSrc)
       return true;

   return false;
}

bool StaticMesh::LoadRenderResourceTexture(DKObject<DKCommandQueue> queue)
{
    if (textureSrc && texture == nullptr)
    {
        DKGraphicsDevice* device = queue->Device();
        DKTextureDescriptor texDesc = {};
        texDesc.textureType = DKTexture::Type2D;
        texDesc.pixelFormat = DKPixelFormat::RGBA8Unorm;
        texDesc.width = textureSrc->Width();
        texDesc.height = textureSrc->Height();
        texDesc.depth = 1;
        texDesc.mipmapLevels = 1;
        texDesc.sampleCount = 1;
        texDesc.arrayLength = 1;
        texDesc.usage = DKTexture::UsageCopyDestination | DKTexture::UsageSampled;
        DKObject<DKTexture> tex = device->CreateTexture(texDesc);
        if (tex)
        {
            size_t bytesPerPixel = textureSrc->BytesPerPixel();
            DKASSERT_DESC(bytesPerPixel
                == DKPixelFormatBytesPerPixel(texDesc.pixelFormat)
                , "BytesPerPixel mismatch!");

            uint32_t width = textureSrc->Width();
            uint32_t height = textureSrc->Height();

            size_t bufferLength = bytesPerPixel * texDesc.width * texDesc.height;
            DKObject<DKGpuBuffer> stagingBuffer = device->CreateBuffer(bufferLength
                , DKGpuBuffer::StorageModeShared, DKCpuCacheModeReadWrite);

            memcpy(stagingBuffer->Contents(), textureSrc->Contents(), bufferLength);
            stagingBuffer->Flush();

            DKObject<DKCommandBuffer> cb = queue->CreateCommandBuffer();
            DKObject<DKCopyCommandEncoder> encoder = cb->CreateCopyCommandEncoder();
            encoder->CopyFromBufferToTexture(stagingBuffer,
                { 0, width, height },
                tex,
                { 0,0, 0,0,0 },
                { width,height,1 });
            encoder->EndEncoding();
            cb->Commit();

            DKLog("Texture created!");
            texture = tex;

            // create sampler
            DKSamplerDescriptor samplerDesc = {};
            samplerDesc.magFilter = DKSamplerDescriptor::MinMagFilterLinear;
            samplerDesc.minFilter = DKSamplerDescriptor::MinMagFilterLinear;
            samplerDesc.addressModeU = DKSamplerDescriptor::AddressModeClampToEdge;
            samplerDesc.addressModeV = DKSamplerDescriptor::AddressModeClampToEdge;
            samplerDesc.addressModeW = DKSamplerDescriptor::AddressModeClampToEdge;
            samplerDesc.maxAnisotropy = 16;

            textureSampler = device->CreateSamplerState(samplerDesc);

            return true;
        }
    }
    return false;
}

bool StaticMesh::LoadShaderFromFile(DKResourcePool& inDKResourcePool
    , const DKString& inFileName
    , DKShaderStage inStage)
{
    DKObject<DKData> data = inDKResourcePool.LoadResourceData(inFileName);
    shaderData[static_cast<int>(inStage)] = data;
    if (data)
        return true;
    return false;
}

bool StaticMesh::LoadRenderResourceShader(DKObject<DKGraphicsDevice> inDevice
    , DKShaderStage inStage)
{
    DKShader shader(shaderData[static_cast<int>(inStage)]);
    DKObject<DKShaderModule> module = inDevice->CreateShaderModule(&shader);
    DKObject<DKShaderFunction> function
        = module->CreateFunction(module->FunctionNames().Value(0));
    shaderModule[static_cast<int>(inStage)] = module;
    shaderFunction[static_cast<int>(inStage)] = function;

    if (module && function)
    {
        if (inStage == DKShaderStage::Vertex)
        {
            DKLog("VertexFunction.VertexAttributes: %d", function->StageInputAttributes().Count());
            for (int i = 0; i < function->StageInputAttributes().Count(); ++i)
            {
                const DKShaderAttribute& attr = function->StageInputAttributes().Value(i);
                DKLog("  --> VertexAttribute[%d]: \"%ls\" (location:%u)", i, (const wchar_t*)attr.name, attr.location);
            }
        }
        return true;
    }

    return false;
}

void StaticMesh::SetupPipelineDecriptor(DKRenderPipelineDescriptor& pipelineDescriptor)
{
    // setup shader
    pipelineDescriptor.vertexFunction = shaderFunction[static_cast<int>(DKShaderStage::Vertex)];
    pipelineDescriptor.fragmentFunction = shaderFunction[static_cast<int>(DKShaderStage::Fragment)];

    uint32_t i = 0;
    for (auto decl : mesh->vertexBuffers[0].declarations)
    {
        pipelineDescriptor.vertexDescriptor.attributes.Add(
            DKVertexAttributeDescriptor{ decl.format, decl.offset, 0, i }
        );
        i++;
    }

    pipelineDescriptor.vertexDescriptor.layouts = {
        { DKVertexStepRate::Vertex, mesh->vertexBuffers[0].size, 0 },
    };
}

void StaticMesh::SetupMaterial(DKObject<DKGraphicsDevice> inDevice)
{
    DKShaderBindingSetLayout layout;
    if (1)
    {
        DKShaderBinding bindings[2] = {
            {
                0,
                DKShader::DescriptorTypeUniformBuffer,
                1,
                nullptr
            }, // Scene Uniform Buffer
            {
                1,
                DKShader::DescriptorTypeTextureSampler,
                1,
                nullptr
            }, // Texture
        };
        layout.bindings.Add(bindings, 2);
    }

    bindSet = inDevice->CreateShaderBindingSet(layout);

    if (bindSet == nullptr)
        return;

    bindSet->SetTexture(1, texture);
    bindSet->SetSamplerState(1, textureSampler);
}

void StaticMesh::SetupExternalUniformBuffer(DKObject<DKGpuBuffer> uniformBuffer, size_t uniformBufferSize, int index)
{
    if (uniformBuffer == nullptr || bindSet == nullptr)
        return;

    bindSet->SetBuffer(index, uniformBuffer, 0, uniformBufferSize);
}

bool StaticMesh::EncodeRenderCommand(DKRenderCommandEncoder* encoder) const
{
    encoder->SetVertexBuffer(mesh->vertexBuffers[0].buffer, 0, 0);
    encoder->SetIndexBuffer(mesh->subMeshes[0].indexBuffer, 0, mesh->subMeshes[0].indexType);
    encoder->SetResources(0, bindSet);
    // draw scene!
    encoder->DrawIndexed(mesh->subMeshes[0].indexCount, 1, 0, 0, 0);
    encoder->EndEncoding();
    return true;
}