#include "model.hpp"
#include "dbc.hpp"
#include <print>

unsigned int texture_from_file(const std::string& directory, const std::string& path);
unsigned int texture_from_internal(const aiScene *scene, size_t index);

unsigned int image_to_texture(unsigned char *data, int width, int height, int nrComponents) {
  unsigned int textureID;
  glGenTextures(1, &textureID);

  GLenum format = GL_RGB;

  switch(nrComponents) {
    case 1:
      format = GL_RED;
      break;
    case 3:
      format = GL_RGB;
      break;
    case 4:
      format = GL_RGBA;
      break;
    default:
      dbc::sentinel($F("Impossible nrComponents={} when loading image", nrComponents));
  }

  glBindTexture(GL_TEXTURE_2D, textureID);
  glTexImage2D(GL_TEXTURE_2D, 0, format, width, height, 0, format, GL_UNSIGNED_BYTE, data);
  glGenerateMipmap(GL_TEXTURE_2D);

  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

  return textureID;
}

unsigned int texture_from_internal(const aiScene *scene, size_t index) {
  const aiTexture *aiTex = scene->mTextures[index];
  unsigned int textureID = 0;
  int width = 0;
  int height = 0;
  int nrComponents = 0;
  unsigned char *data = nullptr;
  const unsigned char* dataBytes = (const unsigned char*)aiTex->pcData;

  if(aiTex->mHeight == 0) {
    // compressed texture (PNG or JPG)
    size_t dataSize = aiTex->mWidth;
    data = stbi_load_from_memory(dataBytes, dataSize, &width, &height, &nrComponents, 0);
    textureID = image_to_texture(data, width, height, nrComponents);
    stbi_image_free(data);
  } else {
    data = (unsigned char*)aiTex->pcData;
    width = aiTex->mWidth;
    height = aiTex->mHeight;
    nrComponents = 4; // either BGRA8888 or RGBA8888
    textureID = image_to_texture(data, width, height, nrComponents);
  }

  return textureID;
}

unsigned int texture_from_file(const std::string& directory, const std::string& path)
{
  std::string filename = directory + "/" + path;

  int width = 0;
  int height = 0;
  int nrComponents = 0;

  unsigned char *data = stbi_load(filename.c_str(), &width, &height, &nrComponents, 0);
  dbc::check(data != nullptr, $F("Failed to load texture {}", filename));

  unsigned int textureID = image_to_texture(data, width, height, nrComponents);

  stbi_image_free(data);

  return textureID;
}

void Model::draw(Shader &shader) {
  for(auto& mesh : meshes) {
    mesh.draw(shader);
  }
}

void Model::load_model() {
  std::string path = directory + "/" + model_path;

  Assimp::Importer importer;

  const aiScene* scene = importer.ReadFile(path,
      aiProcess_Triangulate |
      aiProcess_GenSmoothNormals |
      aiProcess_FlipUVs |
      aiProcess_CalcTangentSpace);

  dbc::check(scene != nullptr, "Assimp ReadFile return null");
  dbc::check(!(scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE), "Assimp says incomplete.");
  dbc::check(scene->mRootNode != nullptr, "Assimp loaded scene doesn't have a root node");

  process_node(scene, scene->mRootNode);
}

void Model::process_node(const aiScene *scene, aiNode *node) {
  for(unsigned int i = 0; i < node->mNumMeshes; i++) {
    aiMesh* mesh = scene->mMeshes[node->mMeshes[i]];
    meshes.push_back(process_mesh(scene, mesh));
  }

  for(unsigned int i = 0; i < node->mNumChildren; i++) {
    process_node(scene, node->mChildren[i]);
  }
}

Mesh Model::process_mesh(const aiScene *scene, aiMesh *mesh) {
  std::vector<Vertex> vertices;
  std::vector<unsigned int> indices;
  std::vector<Texture> textures;

  for(unsigned int i = 0; i < mesh->mNumVertices; i++) {
    glm::vec3 position{
      mesh->mVertices[i].x,
      mesh->mVertices[i].y,
      mesh->mVertices[i].z};

    glm::vec3 normal{};
    glm::vec2 tex_coords{0.0f, 0.0f};

    if(mesh->HasNormals()) {
      normal = glm::vec3(
        mesh->mNormals[i].x,
        mesh->mNormals[i].y,
        mesh->mNormals[i].z);
    }

    if(mesh->mTextureCoords[0]) {
      tex_coords = glm::vec2(
          mesh->mTextureCoords[0][i].x,
          mesh->mTextureCoords[0][i].y
      );
    }

    vertices.emplace_back(position, normal, tex_coords);
  }

  for(unsigned int i = 0; i < mesh->mNumFaces; i++) {
    aiFace& face = mesh->mFaces[i];

    for(unsigned int j = 0; j < face.mNumIndices; j++) {
      indices.emplace_back(face.mIndices[j]);
    }
  }

  aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];

  // we assume a convention for sampler names in the shaders. Each diffuse texture should be named
  // as 'texture_diffuseN' where N is a sequential number ranging from 1 to MAX_SAMPLER_NUMBER.
  // Same applies to other texture as the following list summarizes:
  // diffuse: texture_diffuseN
  // specular: texture_specularN
  // normal: texture_normalN

  // 1. diffuse maps
  load_material_textures(scene, textures, material, aiTextureType_DIFFUSE, "texture_diffuse");

  // 2. specular maps
  load_material_textures(scene, textures, material, aiTextureType_SPECULAR, "texture_specular");

  // 3. normal maps
  load_material_textures(scene, textures, material, aiTextureType_HEIGHT, "texture_normal");

  // 4. height maps
  load_material_textures(scene, textures, material, aiTextureType_AMBIENT, "texture_height");

  return Mesh(vertices, indices, textures);
}

void Model::load_material_textures(const aiScene *scene, std::vector<Texture>& textures, aiMaterial *mat, aiTextureType type, std::string type_name)
{
  for(unsigned int i = 0; i < mat->GetTextureCount(type); i++) {
    aiString str;

    mat->GetTexture(type, i, &str);
    std::string tx_path{str.C_Str()};
    dbc::check(!tx_path.empty(), "Texture has empty path, should be impossible?");

    if(textures_loaded.contains(tx_path)) {
      textures.push_back(textures_loaded.at(tx_path));
    } else {
      unsigned int tx_id = 0;

      // if it's a *# style path then it's internal
      if(tx_path[0] == '*') {
        // get the texture from the internal version
        tx_id = texture_from_internal(scene, std::stoi(tx_path.substr(1)));
      } else {
        // else get it from a file
        tx_id = texture_from_file(directory, tx_path);
      }

      // dubious code here, but seems to be right
      auto& result = textures.emplace_back(tx_id, type_name, tx_path, texture_counts);

      textures_loaded.try_emplace(tx_path, result);
    }
  }
}