#include <stdlib.h>
#include <ply.h>
#include <string.h>
#include <stdio.h>

const char* const ply_formats[] = {
  "ascii",
  "binary_little_endian",
  "binary_big_endian"
};

const char* const ply_header_items[] = {
  "comment",
  "end_header",
  "element",
  "property",
};

const char* const ply_type_strings[] = {
  "char",
  "uchar",
  "short",
  "ushort",
  "int",
  "uint",
  "float",
  "double",
};

const size_t ply_type_sizes[] = {
  1,
  1,
  2,
  2,
  4,
  4,
  4,
  8,
};

PlyPropertyType ply_parse_type(const char* type_string) {
  for(uint32_t i = 0; i < sizeof(ply_type_strings)/sizeof(char*); i++) {
    if(strcmp(ply_type_strings[i], type_string) == 0) {
      return i;
    }
  }
  return PLY_TYPE_INVALID;
}

void ply_free_elements(PlyElement* elements_head) {
  PlyElement* elem = elements_head;
  while(elem != NULL) {
    PlyProperty* prop = elem->properties;
    while(prop != NULL) {
      if(prop->count == PLY_TYPE_INVALID) {
        if(prop->elements != NULL) {
          free(prop->elements);
        }
      } else {
        if(prop->elements != NULL) {
          void** lists = prop->elements;
          for(int i = 0; i < elem->count; i++) {
            if(lists[i] != NULL) {
              free(lists[i]);
            }
          }
          free(prop->elements);
        }

        if(prop->counts != NULL) {
          free(prop->counts);
        }

        if(prop->name != NULL) {
          free(prop->name);
        }
      }
      PlyProperty* last = prop;
      prop = prop->next;
      free(last);
    }

    if(elem->name != NULL) {
      free(elem->name);
    }

    PlyElement* last = elem;
    elem = elem->next;
    free(last);
  }
}

PlyMappings default_ply_mappings = {
  .vertex_element = "vertex",
  .position = {"x", "y", "z"},
  .normal = {"nx", "ny", "nz"},
  .colour = {"red", "green", "blue", "alpha"},
  .texture = {"s", "t"},
  .face_element = "face",
  .index = "vetex_indices",
};

uint16_t ply_conv_index(PlyPropertyType type, void* value, int swap_order) {
  if(type == PLY_TYPE_INVALID) {
    return 0.0f;
  }

  uint8_t* ptr = value;
  uint8_t tmp[8];
  size_t type_size = ply_type_sizes[type];
  for(uint32_t i = 0; i < type_size; i++) {
    if(swap_order) {
      tmp[type_size-i] = ptr[i];
    } else {
      tmp[i] = ptr[i];
    }
  }

  switch(type) {
  case PLY_TYPE_CHAR:
    return (uint16_t)*(int8_t*)(tmp);
  case PLY_TYPE_UCHAR:
    return (uint16_t)*(uint8_t*)(tmp);
  case PLY_TYPE_SHORT:
    return *(int16_t*)(tmp);
  case PLY_TYPE_USHORT:
    return (uint16_t)*(uint16_t*)(tmp);
  case PLY_TYPE_INT:
    return (uint16_t)*(int32_t*)(tmp);
  case PLY_TYPE_UINT:
    return (uint16_t)*(uint32_t*)(tmp);
  default:
    return 0.0f;
  }
}

float ply_conv_value(PlyPropertyType type, void* value, int swap_order) {
  if(type == PLY_TYPE_INVALID) {
    return 0.0f;
  }

  uint8_t* ptr = value;
  uint8_t tmp[8];
  size_t type_size = ply_type_sizes[type];
  for(uint32_t i = 0; i < type_size; i++) {
    if(swap_order) {
      tmp[type_size-i] = ptr[i];
    } else {
      tmp[i] = ptr[i];
    }
  }

  switch(type) {
  case PLY_TYPE_CHAR:
    return (float)*(int8_t*)(tmp)/(float)INT8_MAX;
  case PLY_TYPE_UCHAR:
    return (float)*(uint8_t*)(tmp)/(float)UINT8_MAX;
  case PLY_TYPE_SHORT:
    return (float)*(int16_t*)(tmp)/(float)INT16_MAX;
  case PLY_TYPE_USHORT:
    return (float)*(uint16_t*)(tmp)/(float)UINT16_MAX;
  case PLY_TYPE_INT:
    return (float)*(int32_t*)(tmp)/(float)INT32_MAX;
  case PLY_TYPE_UINT:
    return (float)*(uint32_t*)(tmp)/(float)UINT32_MAX;
  case PLY_TYPE_FLOAT:
    return *(float*)(tmp);
  case PLY_TYPE_DOUBLE:
    return *(float*)(tmp);
  default:
    return 0.0f;
  }
}

int ply_map_values(char* name, int stride, int count, PlyPropertyType type, int mapping_count, char** mapping, void* data, float* out) {
  for(int i = 0; i < mapping_count; i++) {
    if(strcmp(name, mapping[i]) == 0) {
      for(int j = 0; j < count; j++) {
        size_t type_size = ply_type_sizes[type];
        out[(stride*j)+i] = ply_conv_value(type, data + (type_size*j), 0);
      }
      return 1;
    }
  }
  return 0;
}

PlyMesh ply_load_mesh(char* filename, PlyMappings mappings) {
  PlyMesh mesh = {
    .vertex_count = 0,
    .position = NULL,
    .colour = NULL,
    .uv = NULL,
    .normal = NULL,

    .index_count = 0,
    .index = NULL,
  };

  FILE* ply = fopen(filename, "r");
  if(ply == NULL) {
    fprintf(stderr, "file not found %s\n", filename);
    return mesh;
  }

  uint8_t buffer[1024] = {0};
  char magic[] = {'p', 'l', 'y', '\n'};
  int read = fread(buffer, 1, sizeof(magic), ply);
  if(read != sizeof(magic)) {
    fprintf(stderr, "magic not read\n");
    goto close;
  } else if (memcmp(buffer, magic, sizeof(magic)) != 0) {
    fprintf(stderr, "magic not match\n");
    goto close;
  }

  char* line = fgets((char*)buffer, 1024, ply);
  if(line == NULL) {
    fprintf(stderr, "file read err\n");
    goto close;
  }

  char* format_str = "format ";
  if(memcmp(format_str, line, 7) != 0) {
    fprintf(stderr, "format start not found\n");
    goto close;
  }

  char* format_start = &line[7];
  PlyFormatType format_id = PLY_FORMAT_INVALID;
  for(uint32_t i = 0; i < sizeof(ply_formats)/sizeof(char*); i++) {
    if(memcmp(format_start, ply_formats[i], strlen(ply_formats[i])) == 0) {
      format_id = i;
      break;
    }
  }

  if(format_id != PLY_FORMAT_BE && format_id != PLY_FORMAT_LE) {
    fprintf(stderr, "can only parse binary PLY files\n");
    goto close;
  }

  PlyElement* elements_head = NULL;
  PlyElement* element = NULL;

  while(1) {
    char* line = fgets((char*)buffer, 1024, ply);
    PlyHeaderType header_id = PLY_HEADER_INVALID;
    for(uint32_t i = 0; i < sizeof(ply_header_items)/sizeof(char*); i++) {
      if(memcmp(line, ply_header_items[i], strlen(ply_header_items[i])) == 0) {
        header_id = i;
        break;
      }
    }

    if(header_id == PLY_FORMAT_INVALID) {
      fprintf(stderr, "unknown header line: %s\n", line);
      goto clean;
    } else if (header_id == PLY_HEADER_COMMENT) {
      continue;
    } else if (header_id == PLY_HEADER_END) {
      break;
    } else if (header_id == PLY_HEADER_ELEMENT) {
      PlyElement* new_element = malloc(sizeof(PlyElement));
      new_element->next = NULL;
      new_element->properties = NULL;

      char* first = strtok((char*)buffer, " ");
      (void)first;
      char* name = strtok(NULL, " ");
      char* count = strtok(NULL, " ");
      if(name == NULL || count == NULL) {
        fprintf(stderr, "BAD_ELEMENT: %s\n", line);
        goto clean;
      }
      unsigned long name_len = strlen(name) + 1;
      new_element->name = malloc(name_len);
      memcpy(new_element->name, name, name_len);

      new_element->count = atoi(count);

      if(elements_head == NULL) {
        elements_head = new_element;
        element = new_element;
      } else {
        element->next = new_element;
        element = new_element;
      }
    } else if (header_id == PLY_HEADER_PROPERTY) {
      if(element == NULL) {
        fprintf(stderr, "parsed property before any element\n");
        goto clean;
      }
      PlyProperty* new_property = malloc(sizeof(PlyProperty));
      new_property->name = NULL;
      new_property->count = PLY_TYPE_INVALID;
      new_property->element = PLY_TYPE_INVALID;
      new_property->next = NULL;
      new_property->elements = NULL;
      new_property->counts = NULL;

      char* extra = strtok((char*)buffer, " ");
      (void)extra;

      char* type_str = strtok(NULL, " ");
      if(strcmp(type_str, "list") == 0) {
        type_str = strtok(NULL, " ");
        new_property->count = ply_parse_type(type_str);
        type_str = strtok(NULL, " ");
      }

      new_property->element = ply_parse_type(type_str);

      type_str = strtok(NULL, " ");
      char* name = strtok(type_str, "\n");
      if(name == NULL) {
        fprintf(stderr, "invalid property\n");
        goto clean;
      }

      unsigned long name_len = strlen(name) + 1;
      new_property->name = malloc(name_len);
      memcpy(new_property->name, name, name_len);

      PlyProperty* cur = element->properties;
      if(cur == NULL) {
        element->properties = new_property;
      } else {
        while(cur->next != NULL) {
          cur = cur->next;
        }
        cur->next = new_property;
      }
    }

    if(line == NULL) {
      goto clean;
    }
  }

  PlyElement* elem = elements_head;
  while(elem != NULL) {
    PlyProperty* prop = elem->properties;
    while(prop != NULL) {
      if(prop->count == PLY_TYPE_INVALID) {
        prop->elements = malloc(ply_type_sizes[prop->element]*elem->count);
        prop->counts = NULL;
      } else {
        prop->elements = malloc(sizeof(void*)*elem->count);
        prop->counts = malloc(sizeof(ply_type_sizes[prop->count])*elem->count);
      }
      prop = prop->next;
    }

    for(int i = 0; i < elem->count; i++) {
      PlyProperty* prop = elem->properties;
      while(prop != NULL) {
        if(prop->count == PLY_TYPE_INVALID) {
          int read = fread(prop->elements + (i*ply_type_sizes[prop->element]), ply_type_sizes[prop->element], 1, ply);
          if(read != 1) {
            fprintf(stderr, "prop read error\n");
            goto clean;
          }
        } else {
          uint64_t size = 0;
          uint64_t read = fread(&size, ply_type_sizes[prop->count], 1, ply);
          if(read != 1) {
            fprintf(stderr, "prop read error\n");
            goto clean;
          }

          memcpy(prop->counts + (i*ply_type_sizes[prop->count]), &size, ply_type_sizes[prop->count]);
          void* elements = malloc(ply_type_sizes[prop->element]*size);
          read = fread(elements, ply_type_sizes[prop->element], size, ply);
          if(read != size) {
            fprintf(stderr, "prop read error\n");
            goto clean;
          }
          memcpy(prop->elements + (i*sizeof(void*)), &elements, sizeof(void*));
        }
        prop = prop->next;
      }
    }
    elem = elem->next;
  }

  elem = elements_head;
  while(elem != NULL) {
    int type = -1;
    if(strcmp(elem->name, mappings.vertex_element) == 0) {
      type = 0;
      mesh.position = malloc(sizeof(vec3)*elem->count);
      mesh.colour = malloc(sizeof(vec4)*elem->count);
      mesh.normal = malloc(sizeof(vec3)*elem->count);
      mesh.uv = malloc(sizeof(vec2)*elem->count);
      mesh.vertex_count = elem->count;
    } else if(strcmp(elem->name, mappings.face_element) == 0) {
      type = 1;
      mesh.index_count = 3*elem->count;
      mesh.index = malloc(3*elem->count*sizeof(uint16_t));
    } else {
      continue;
    }

    PlyProperty* prop = elem->properties;
    while(prop != NULL) {
      if(prop->count == PLY_TYPE_INVALID && type == 0) {
        if(ply_map_values(prop->name,
              3,
              elem->count,
              prop->element,
              sizeof(mappings.position)/sizeof(char*),
              mappings.position,
              prop->elements,
              (float*)mesh.position)) {
        } else if(ply_map_values(prop->name,
              4,
              elem->count,
              prop->element,
              sizeof(mappings.colour)/sizeof(char*),
              mappings.colour,
              prop->elements,
              (float*)mesh.colour)) {
        } else if(ply_map_values(prop->name,
              3,
              elem->count,
              prop->element,
              sizeof(mappings.normal)/sizeof(char*),
              mappings.normal,
              prop->elements,
              (float*)mesh.normal)) {
        } else if(ply_map_values(prop->name,
              2,
              elem->count,
              prop->element,
              sizeof(mappings.texture)/sizeof(char*),
              mappings.texture,
              prop->elements,
              (float*)mesh.uv)) {
        }
      } else if (prop->count != PLY_TYPE_INVALID && type == 1) {
        for(int i = 0; i < elem->count; i++) {
          int type_size = ply_type_sizes[prop->element];
          void ** lists = prop->elements;
          mesh.index[i*3 + 0] = ply_conv_index(prop->element, lists[i] + (0*type_size), 0);
          mesh.index[i*3 + 1] = ply_conv_index(prop->element, lists[i] + (1*type_size), 0);
          mesh.index[i*3 + 2] = ply_conv_index(prop->element, lists[i] + (2*type_size), 0);
        }
      }
      prop = prop->next;
    }
    elem = elem->next;
  }

clean:
  ply_free_elements(elements_head);
close:
  fclose(ply);
  return mesh;
}