// Copyright (c) 2023 Dominic Masters
// 
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#include "display/shader/SimpleTexturedShader.hpp"

#include <fstream>
#include "slang.h"
#include "slang-gfx.h"
using namespace slang;

using namespace Dawn;

void SimpleTexturedShader::getStages(
  const enum ShaderOpenGLVariant variant,
  const struct SimpleTexturedShaderData *rel,
  std::vector<std::shared_ptr<ShaderStage>> &stages,
  std::vector<struct ShaderParameter> &parameters,
  std::vector<struct IShaderStructure> &structures
) {
  // Stages
  std::shared_ptr<ShaderStage> vertex;
  std::shared_ptr<ShaderStage> fragment;

  std::string shader = R"(
    cbuffer Uniforms {
      float4x4 u_Projection;
      float4x4 u_View;
      float4x4 u_Model;
      float4 u_Color;
      bool u_HasTexture;
      uniform Sampler2D u_Texture;
    };

    struct AssembledVertex {
      float3 position : POSITION;
      float2 texcoord : TEXCOORD;
    };

    struct Fragment {
      float4 color;
    };

    struct VertexStageOutput {
      float2 uv : UV;
      float4 sv_position : SV_Position;
    };

    [shader("vertex")]
    VertexStageOutput vertexMain(
      AssembledVertex assembledVertex
    ) {
      VertexStageOutput output;

      float3 position = assembledVertex.position;

      output.uv = assembledVertex.texcoord;

      output.sv_position = mul(
        float4(position, 1.0),
        mul(u_Model, mul(u_View, u_Projection))
      );

      return output;
    }

    [shader("fragment")]
    Fragment fragmentMain(
      float2 uv: UV
    ) : SV_Target {
      Fragment output;
      if(u_HasTexture) {
        output.color = u_Texture.Sample(uv) * u_Color;
      } else {
        output.color = u_Color;
      }
      return output;
    }
  )";
  Slang::ComPtr<IGlobalSession> globalSession;
  createGlobalSession(globalSession.writeRef());

  SessionDesc sessionDesc;

  TargetDesc targetDesc;
  targetDesc.format = SLANG_GLSL;
  targetDesc.profile = globalSession->findProfile("glsl_330");
  sessionDesc.targets = &targetDesc;
  sessionDesc.targetCount = 1;

  Slang::ComPtr<IBlob> diagnostics;
  const char* searchPaths[] = { "/home/yourwishes/htdocs/Dawn/assets/shaders/" };
  sessionDesc.searchPaths = searchPaths;
  sessionDesc.searchPathCount = 1;

  Slang::ComPtr<ISession> session;
  globalSession->createSession(sessionDesc, session.writeRef());
  auto module = session->loadModuleFromSourceString(
    "hello-world.slang",
    "hello-world.slang",
    shader.c_str(),
    diagnostics.writeRef()
  );
  if(diagnostics) {
    assertUnreachable("Failed to load module %s", (const char*) diagnostics->getBufferPointer());
    return;
  }

  Slang::ComPtr<IEntryPoint> vertexEntryPoint;
  Slang::ComPtr<IEntryPoint> fragEntryPoint;
  module->findEntryPointByName("vertexMain", vertexEntryPoint.writeRef());
  module->findEntryPointByName("fragmentMain", fragEntryPoint.writeRef());

  IComponentType* components[] = { module, vertexEntryPoint, fragEntryPoint };
  Slang::ComPtr<IComponentType> program;
  session->createCompositeComponentType(
    components,
    sizeof(components) / sizeof(components[0]),
    program.writeRef()
  );

  Slang::ComPtr<IComponentType> linkedProgram;
  auto result = program->link(linkedProgram.writeRef(), diagnostics.writeRef());
  std::cout << "Result: " << result << std::endl;
  if(diagnostics) {
    assertUnreachable("%s\n", (const char*) diagnostics->getBufferPointer());
    return;
  }

  int entryPointIndex = 0;
  int targetIndex = 0; // only one target
  Slang::ComPtr<IBlob> vertexBlob;
  result = linkedProgram->getEntryPointCode(
    entryPointIndex,
    targetIndex,
    vertexBlob.writeRef(),
    diagnostics.writeRef()
  );
  if(diagnostics) {
    assertUnreachable("%s\n", (const char*) diagnostics->getBufferPointer());
    return;
  }

  slang::ProgramLayout* layout = program->getLayout();
  unsigned parameterCount = layout->getParameterCount();
  for(unsigned pp = 0; pp < parameterCount; pp++) {
    slang::VariableLayoutReflection* parameter = layout->getParameterByIndex(pp);
    std::cout << "Parameter: " << parameter->getName() << std::endl;
    
    auto layout = parameter->getTypeLayout();
    auto fields = layout->getFieldCount();
    for(unsigned ff = 0; ff < fields; ff++) {
      slang::VariableLayoutReflection* field = layout->getFieldByIndex(ff);
      std::string fieldName = field->getName();
      std::cout << "Field: " << fieldName << std::endl;
    }
  }

  std::string vertexString = (const char*)vertexBlob->getBufferPointer();

  entryPointIndex = 1;
  Slang::ComPtr<IBlob> fragmentBlob;
  result = linkedProgram->getEntryPointCode(
    entryPointIndex,
    targetIndex,
    fragmentBlob.writeRef(),
    diagnostics.writeRef()
  );
  if(diagnostics) {
    assertUnreachable("%s\n", (const char*) diagnostics->getBufferPointer());
    return;
  }

  std::string fragmentString = (const char*)fragmentBlob->getBufferPointer();

  vertex = std::make_shared<ShaderStage>(
    ShaderStageType::VERTEX, vertexString
  );
  stages.push_back(vertex);

  fragment = std::make_shared<ShaderStage>(
    ShaderStageType::FRAGMENT, fragmentString
  );
  stages.push_back(fragment);

  structures.push_back(ShaderStructure<struct SimpleTexturedShaderDataSub>(
    "block_SLANG_ParameterGroup_Uniforms_std140_0",
    &rel->data,
    ShaderOpenGLStructureType::STD140,
    [](const SimpleTexturedShaderDataSub &data, std::vector<struct ShaderParameter> &parameters) {
      parameters.push_back(ShaderParameter(
        "u_Projection_0",
        &data.projection,
        ShaderParameterType::MAT4
      ));

      parameters.push_back(ShaderParameter(
        "u_View_0",
        &data.view,
        ShaderParameterType::MAT4
      ));

      parameters.push_back(ShaderParameter(
        "u_Model_0",
        &data.model,
        ShaderParameterType::MAT4
      ));

      parameters.push_back(ShaderParameter(
        "u_Color_0",
        &data.color,
        ShaderParameterType::COLOR
      ));

      parameters.push_back(ShaderParameter(
        "u_HasTexture_0",
        &data.hasTexture,
        ShaderParameterType::BOOLEAN
      ));
    }
  ));

  // Parameters
  parameters.push_back(ShaderParameter(
    "Uniforms_u_Texture_0",
    &rel->texture,
    ShaderParameterType::TEXTURE
  ));
}