// 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;
    };

    layout(binding = 1)
    uniform Sampler2D u_Texture;

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

    struct CoarseVertex {
    };

    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::ProgramLayout* layout = program->getLayout();

  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;
  }

  std::string vertexString = (const char*)vertexBlob->getBufferPointer();
  std::ofstream out("/home/yourwishes/htdocs/Dawn/vertex.glsl");
  out << vertexString;
  out.close();

  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();
  out.open("/home/yourwishes/htdocs/Dawn/fragment.glsl");
  out << fragmentString;
  out.close();

  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(
  //   "u_Projection",
  //   &rel->projection,
  //   ShaderParameterType::MAT4
  // ));

  // parameters.push_back(ShaderParameter(
  //   "u_View",
  //   &rel->view,
  //   ShaderParameterType::MAT4
  // ));

  // parameters.push_back(ShaderParameter(
  //   "u_Model",
  //   &rel->model,
  //   ShaderParameterType::MAT4
  // ));

  // parameters.push_back(ShaderParameter(
  //   "u_Color",
  //   &rel->color,
  //   ShaderParameterType::COLOR
  // ));

  // parameters.push_back(ShaderParameter(
  //   "u_HasTexture",
  //   &rel->hasTexture,
  //   ShaderParameterType::BOOLEAN
  // ));

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