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

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_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  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, bool with_textures) {
  for(auto& mesh : meshes) {
    mesh.draw(shader, with_textures);
  }
}

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

  fmt::println("Loading 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]);
  }
}

void Model::SetVertexBoneData(Vertex& vertex, int boneID, float weight)
{
  // BUG: this will ignore bones after the max, but it seems to always be 4, so is there an off-by-one?
  if(vertex.bone_count < MAX_BONE_INFLUENCE) {
    vertex.bone_ids[vertex.bone_count] = boneID;
    vertex.weights[vertex.bone_count] = weight;
    vertex.bone_count++;
  }
}

void Model::ExtractBoneWeightForVertices(std::vector<Vertex>& vertices, aiMesh *mesh, const aiScene* scene) {

  for(size_t boneIndex = 0; boneIndex < mesh->mNumBones; ++boneIndex) {
    int boneID = -1;
    std::string bone_name = mesh->mBones[boneIndex]->mName.C_Str();

    if(bone_info_map.contains(bone_name)) {
      boneID = bone_info_map.at(bone_name).id;
    } else {
      fmt::println("FOUND BONE: {}", bone_name);
      bone_info_map.try_emplace(bone_name,
          bone_counter,
          AssimpGLMHelpers::ConvertMatrixToGLMFormat(mesh->mBones[boneIndex]->mOffsetMatrix));
      boneID = bone_counter;
      bone_counter++;
    }

    dbc::check(boneID != -1, "There's no bones!");
    auto weights = mesh->mBones[boneIndex]->mWeights;
    int numWeights = mesh->mBones[boneIndex]->mNumWeights;

    for(int weightIndex = 0; weightIndex < numWeights; ++weightIndex) {
      int vertexId = weights[weightIndex].mVertexId;
      float weight = weights[weightIndex].mWeight;
      dbc::check((size_t)vertexId <= vertices.size(), "vertexID is too large");
      SetVertexBoneData(vertices[vertexId], boneID, weight);
    }
  }
}

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};
    glm::vec3 tangent;
    glm::vec3 bitangent;

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

    // if has height/normal map
    if(mesh->HasTangentsAndBitangents()) {
      tangent = glm::vec3(
          mesh->mTangents[i].x,
          mesh->mTangents[i].y,
          mesh->mTangents[i].z
      );

      bitangent = glm::vec3(
          mesh->mBitangents[i].x,
          mesh->mBitangents[i].y,
          mesh->mBitangents[i].z
      );
    }

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

  ExtractBoneWeightForVertices(vertices, mesh, scene);

  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, possibly wrong?
  load_material_textures(scene, textures, material, aiTextureType_NORMALS, "texture_normal");

  // 4. height maps
  load_material_textures(scene, textures, material, aiTextureType_HEIGHT, "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);
    }
  }
}