#include "DirectXMeshPainter.h"

#include "DrawingContext.h"
#include "DrawingSurface.h"
#include "Scene.h"
#include "SceneModel.h"

#include "TriMesh.h"
#include "LineMesh.h"

#include "DirectXShaderProgram.h"
#include "DirectXMesh.h"
#include "DirectXInterface.h"

#include <windows.h>
#include <iostream>
#include <filesystem>

#include <directx/d3dx12.h>

DirectXMeshPainter::DirectXMeshPainter()
{
    auto shader_path = std::filesystem::path(__FILE__).parent_path() / "shaders.hlsl";
    mShaderProgram = std::make_unique<DirectXShaderProgram>(shader_path, shader_path);
}

void DirectXMeshPainter::setDxInterface(DirectXInterface* dxInterface)
{
    mDxInterface = dxInterface;
}

void DirectXMeshPainter::initializeShader()
{

}

void DirectXMeshPainter::initializeD3d()
{
    createRootSignature();
    createPipelineStateObject();
}

void DirectXMeshPainter::createRootSignature()
{
    CD3DX12_ROOT_SIGNATURE_DESC rootSignatureDesc;
    rootSignatureDesc.Init(0, nullptr, 0, nullptr, D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

    Microsoft::WRL::ComPtr<ID3DBlob> signature;
    Microsoft::WRL::ComPtr<ID3DBlob> error;
    D3D12SerializeRootSignature(&rootSignatureDesc, D3D_ROOT_SIGNATURE_VERSION_1, &signature, &error);
    mDxInterface->getD3dDevice()->CreateRootSignature(0, signature->GetBufferPointer(), signature->GetBufferSize(), IID_PPV_ARGS(&mRootSignature));
}

void DirectXMeshPainter::createPipelineStateObject()
{
    auto vert_shader = mShaderProgram->getVertexShader();
    auto pixel_shader = mShaderProgram->getPixelShader();

    D3D12_INPUT_ELEMENT_DESC inputElementDescs[] =
    {
        { "POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 0, D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA, 0 },
        { "COLOR", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0, 12, D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA, 0 }
    };

    D3D12_GRAPHICS_PIPELINE_STATE_DESC psoDesc = {};
    psoDesc.InputLayout = { inputElementDescs, _countof(inputElementDescs) };
    psoDesc.pRootSignature = mRootSignature.Get();
    psoDesc.VS = CD3DX12_SHADER_BYTECODE(vert_shader);
    psoDesc.PS = CD3DX12_SHADER_BYTECODE(pixel_shader);
    psoDesc.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT);
    psoDesc.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT);
    psoDesc.DepthStencilState.DepthEnable = FALSE;
    psoDesc.DepthStencilState.StencilEnable = FALSE;
    psoDesc.SampleMask = UINT_MAX;
    psoDesc.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
    psoDesc.NumRenderTargets = 1;
    psoDesc.RTVFormats[0] = DXGI_FORMAT_R8G8B8A8_UNORM;
    psoDesc.SampleDesc.Count = 1;
    mDxInterface->getD3dDevice()->CreateGraphicsPipelineState(&psoDesc, IID_PPV_ARGS(&mPipelineState));
}

ID3D12PipelineState* DirectXMeshPainter::getPipelineState() const
{
    return mPipelineState.Get();
}

ID3D12RootSignature* DirectXMeshPainter::getRootSignature() const
{
    return mRootSignature.Get();
}

void DirectXMeshPainter::updateCommandList(ID3D12GraphicsCommandList* commandList)
{
    commandList->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);

    for (const auto& mesh : mDxMeshes)
    {
        commandList->IASetIndexBuffer(&mesh->getIndexBufferView());
        commandList->IASetVertexBuffers(0, 1, &mesh->getVertBufferView());   
        commandList->DrawInstanced(mesh->getVertBufferView().SizeInBytes/ mesh->getVertBufferView().StrideInBytes, 1, 0, 0);
        //commandList->DrawIndexedInstanced(3, 1, 0, 0, 0);
    }
}

void DirectXMeshPainter::updateBuffers(DrawingContext* context)
{
    mDxMeshes.clear();
    auto scene = context->getSurface()->getScene();
    for (const auto item : scene->getItems())
    {
        if (!item->isVisible())
        {
            continue;
        }

        if (item->getType() == SceneItem::Type::MODEL)
        {
            auto model = dynamic_cast<SceneModel*>(item);
            if (model->getGeometry())
            {
                continue;
            }

            auto dx_mesh = std::make_unique<DirectXMesh>(model);
            dx_mesh->update(context, mDxInterface->getD3dDevice());
            mDxMeshes.push_back(std::move(dx_mesh));
        }
    }
}


