#include "Shader.h"

#include <d3d11.h>
#include <d3d12.h>
#include <iostream>

#include "DXBuffer.h"
#include "File.h"
#include "Text.h"

using namespace Framework;

// Contents of the Shader class

// Constructor
Shader::Shader()
    : ReferenceCounter()
{
    type = UNBEKANNT;
    constBuffers = new RCArray<DXBuffer>();
}

// Destructor
Shader::~Shader()
{
    constBuffers->release();
}

// Deletes a constant buffer
//  index: the index of the buffer to be deleted. Buffer 0 cannot
//  be deleted while buffer 1 still exists, etc.
bool Shader::removeConstBuffer(int index)
{
    if (index < 0) return 0;
    bool ok = 1;
    constBuffers->set(0, index);
    return 1;
}

// Copies data into a constant buffer
//  zD3d11Context: The context object used for copying
//  data: A pointer to a byte array the size of the buffer
//  index: The index of the buffer
//  laen: The length of the data in bytes (-1 for the maximum size of the
//  buffer)
bool Shader::fillConstBuffer(char* data, int index, int len)
{
    if (index < 0 || index > constBuffers->getLastIndex()) return 0;
    DXBuffer* zB = constBuffers->z(index);
    if (!zB) return 0;
    if (len < 0) len = zB->getElementCount() * zB->getElementLength();
    zB->setData(data);
    zB->copyToGPU(len);
    return 1;
}

// Returns the length of a constant buffer
//  index: The index of the buffer
int Shader::getConstBufferLaenge(int index) const
{
    if (index < 0 || index > constBuffers->getLastIndex()) return 0;
    DXBuffer* zB = constBuffers->z(index);
    if (!zB) return 0;
    return zB->getElementCount() * zB->getElementLength();
}

// Returns the shader type
ShaderType Shader::getType() const
{
    return type;
}

//! Returns the index of the first uninitialized buffer
int Shader::getFirstUninitializedBufferIndex() const
{
    for (int index = 0; index < constBuffers->getEntryCount(); index++)
    {
        if (!constBuffers->has(index) || !constBuffers->z(index)) return index;
    }
    return constBuffers->getEntryCount();
}

DX11Shader::DX11Shader(
    ID3D11Device* device, ID3D11DeviceContext* context, Critical& deviceLock)
    : Shader(),
      deviceLock(deviceLock)
{
    this->device = device;
    this->context = context;
}

DX11Shader::~DX11Shader() {}

// Creates a constant buffer that passes constant data to the shader
// A maximum of 14 buffers can be created
//  zD3d11Device: The device used to create the buffer
//  groesse: The size of the buffer in bytes
//  index: The position of the buffer in the buffer array. Existing buffer
//  is replaced. Buffer 1 cannot be created if buffer 0 has not yet
//  been created, etc.
bool DX11Shader::createConstBuffer(int groesse, int index)
{
    if (index < 0 || index >= 14) return 0;
    bool ok = 1;
    while ((groesse / 16) * 16
           != groesse) // only multiples of 16 are allowed as size
        groesse++;
    while (!constBuffers->has(index))
        constBuffers->add(0);
    constBuffers->set(
        new DX11Buffer(
            1, device, context, D3D11_BIND_CONSTANT_BUFFER, deviceLock),
        index);
    constBuffers->z(index)->setLength(groesse);
    return 1;
}

// Contents of the PixelShader class

// Constructor
DX11PixelShader::DX11PixelShader(
    ID3D11Device* device, ID3D11DeviceContext* context, Critical& deviceLock)
    : DX11Shader(device, context, deviceLock)
{
    pixelShader = 0;
}

// Destructor
DX11PixelShader::~DX11PixelShader()
{
    if (pixelShader) pixelShader->Release();
}

// Sets the compiled shader
//  bytes: The bytes of the compiled code
//  length: the length of the byte array
//  return: true if bytes is valid, false otherwise
bool DX11PixelShader::setCompiledByteArray(unsigned char* bytes, int length)
{
    deviceLock.lock();
    HRESULT result = device->CreatePixelShader(bytes, length, 0, &pixelShader);
    deviceLock.unlock();
    return result == S_OK;
}

// After calling this function, this shader is used as pixel shader
//  zD3d11Context: The context object used with the shader
void DX11PixelShader::useShader()
{
    int maxI = constBuffers->getLastIndex();
    for (int i = 0; i <= maxI; i++)
    {
        if (!constBuffers->z(i)) continue;
        if (!((DX11Buffer*)constBuffers->z(i))->zBuffer())
            constBuffers->z(i)->copyToGPU();
        ID3D11Buffer* buf = ((DX11Buffer*)constBuffers->z(i))->zBuffer();
        context->PSSetConstantBuffers(i, 1, &buf);
    }
    if (pixelShader) context->PSSetShader(pixelShader, 0, 0);
}

// Contents of the VertexShader class

// Constructor
DX11VertexShader::DX11VertexShader(
    ID3D11Device* device, ID3D11DeviceContext* context, Critical& deviceLock)
    : DX11Shader(device, context, deviceLock)
{
    vertexShader = 0;
    inputLayout = 0;
    shaderByteBuffer = 0;
    byteBufferSize = 0;
}

// Destructor
DX11VertexShader::~DX11VertexShader()
{
    if (vertexShader) vertexShader->Release();
    if (inputLayout) inputLayout->Release();
}

// Sets the compiled shader
//  bytes: The bytes of the compiled code
//  length: the length of the byte array
//  return: true if bytes is valid, false otherwise
bool DX11VertexShader::setCompiledByteArray(unsigned char* bytes, int length)
{
    shaderByteBuffer = (unsigned char*)bytes;
    byteBufferSize = length;
    deviceLock.lock();
    HRESULT result
        = device->CreateVertexShader(bytes, length, 0, &vertexShader);
    deviceLock.unlock();
    return result == S_OK;
}

// Creates an InputLayout for the shader
// Must only be called after compile
//  zD3d11Device: The device used to create the layout
//  descArray: An array with initialization data
//  anz: The number of elements in the array
bool DX11VertexShader::createInputLayout(
    D3D11_INPUT_ELEMENT_DESC* descArray, int anz)
{
    if (!shaderByteBuffer) return 0;
    if (inputLayout) inputLayout->Release();
    inputLayout = 0;
    deviceLock.lock();
    HRESULT res = device->CreateInputLayout(
        descArray, anz, shaderByteBuffer, byteBufferSize, &inputLayout);
    deviceLock.unlock();
    if (res == S_OK)
    {
        shaderByteBuffer = 0;
        byteBufferSize = 0;
    }
    return res == S_OK;
}

// After calling this function, this shader is used as vertex shader
//  zD3d11Context: The context object used with the shader
void DX11VertexShader::useShader()
{
    int maxI = constBuffers->getLastIndex();
    for (int i = 0; i <= maxI; i++)
    {
        if (!constBuffers->z(i)) continue;
        if (!((DX11Buffer*)constBuffers->z(i))->zBuffer())
            constBuffers->z(i)->copyToGPU();
        ID3D11Buffer* buf = ((DX11Buffer*)constBuffers->z(i))->zBuffer();
        context->VSSetConstantBuffers(i, 1, &buf);
    }
    if (inputLayout) context->IASetInputLayout(inputLayout);
    if (vertexShader) context->VSSetShader(vertexShader, 0, 0);
}