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

PlyMesh ply_load_mesh(char* filename) {
  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;
  }

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